{"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"},{"sourceId":13665986,"sourceType":"datasetVersion","datasetId":8688768},{"sourceId":13847805,"sourceType":"datasetVersion","datasetId":8820508},{"sourceId":14519387,"sourceType":"datasetVersion","datasetId":8815386}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# !pip install numpy","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-27T17:06:59.326836Z","iopub.execute_input":"2025-12-27T17:06:59.327146Z","iopub.status.idle":"2025-12-27T17:06:59.333632Z","shell.execute_reply.started":"2025-12-27T17:06:59.327123Z","shell.execute_reply":"2025-12-27T17:06:59.332556Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport shutil\nimport subprocess\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport tempfile\nimport warnings\n\nwarnings.filterwarnings('ignore')\n\n# ==================================================================================\n# 1. CONFIGURATION\n# ==================================================================================\nclass TaxConfig:\n    RAW_DATA_DIR = '/kaggle/input/cafa-6-protein-function-prediction/'\n    IDS_INPUT_DIR = '/kaggle/input/datacafa' \n    OUTPUT_DIR = '/kaggle/working/'\n    \n    TAXONKIT_BIN = \"./taxonkit\"\n    DATA_DIR = \"./tax_data\"\n    TAXONKIT_VERSION = \"v0.20.0\"\n    \n    # Features\n    RANK_NAMES = ['kingdom', 'phylum', 'class', 'order', 'family', 'genus', 'species']\n    UNKNOWN_TOKEN = \"<UNK>\"\n    ROOT_TAXID = \"1\"\n    BATCH_SIZE = 50000\n\ncfg = TaxConfig()\n\n# ==================================================================================\n# 2. SETUP TOOLS\n# ==================================================================================\ndef setup_taxonkit():\n    if not os.path.exists(cfg.TAXONKIT_BIN):\n        print(\"⚙️ Installing TaxonKit...\")\n        url = f\"https://github.com/shenwei356/taxonkit/releases/download/{cfg.TAXONKIT_VERSION}/taxonkit_linux_amd64.tar.gz\"\n        subprocess.run(f\"wget -q {url} -O taxonkit.tar.gz\", shell=True, check=True)\n        subprocess.run(\"tar -xzf taxonkit.tar.gz taxonkit\", shell=True, check=True)\n        subprocess.run(\"chmod +x taxonkit\", shell=True, check=True)\n        os.remove(\"taxonkit.tar.gz\")\n\n    os.makedirs(cfg.DATA_DIR, exist_ok=True)\n    required = [\"nodes.dmp\", \"names.dmp\", \"merged.dmp\", \"delnodes.dmp\"]\n    if not all(os.path.exists(os.path.join(cfg.DATA_DIR, f)) for f in required):\n        print(\"⚙️ Downloading NCBI Taxonomy Database...\")\n        subprocess.run(\"wget -q https://ftp.ncbi.nlm.nih.gov/pub/taxonomy/taxdump.tar.gz\", shell=True, check=True)\n        subprocess.run(f\"tar -xzf taxdump.tar.gz -C {cfg.DATA_DIR} \" + \" \".join(required), shell=True, check=True)\n        os.remove(\"taxdump.tar.gz\")\n    print(\"✅ TaxonKit Ready.\")\n\n# ==================================================================================\n# 3. CORRECTED PARSING LOGIC\n# ==================================================================================\ndef extract_taxids_from_fasta(fasta_path):\n    \"\"\"Handles both UniProt (OX=) and CAFA (Space-separated) formats.\"\"\"\n    if not os.path.exists(fasta_path):\n        print(f\"⚠️ Warning: {fasta_path} not found.\")\n        return {}\n\n    pid_to_taxid = {}\n    with open(fasta_path, 'r') as f:\n        for line in tqdm(f, desc=f\"Scanning {os.path.basename(fasta_path)}\"):\n            if line.startswith('>'):\n                parts = line.strip().split()\n                if not parts: continue\n                \n                # 1. Extract Protein ID\n                raw_id = parts[0][1:] # Remove >\n                pid = raw_id.split('|')[1] if '|' in raw_id else raw_id\n                \n                # 2. Extract TaxID (Hybrid Logic)\n                taxid = cfg.ROOT_TAXID\n                \n                # Check for CAFA Test Format: >ID TaxID (e.g. >A0A0C5B5G6 9606)\n                if len(parts) > 1 and parts[1].isdigit():\n                    taxid = parts[1]\n                else:\n                    # Check for UniProt Train Format: OX=1234\n                    for p in parts:\n                        if p.startswith('OX='):\n                            taxid = ''.join(filter(str.isdigit, p[3:]))\n                            break\n                \n                pid_to_taxid[pid] = taxid\n    return pid_to_taxid\n\ndef load_clean_ids(path):\n    ids = np.load(path, allow_pickle=True)\n    return [str(x).split('|')[1] if '|' in str(x) else str(x) for x in ids]\n\n# ==================================================================================\n# 4. TAXONKIT QUERY\n# ==================================================================================\ndef get_lineages_reformat2(taxids):\n    unique_taxids = sorted(list(set(taxids)))\n    results = {}\n    \n    # STEP 1: Update merged/old IDs using 'lineage -c'\n    # This replaces the \"Filter\" step you saw in the docs\n    print(\"🔄 Updating merged/old TaxIDs...\")\n    id_map = {} # {old_id: new_id}\n    with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmp:\n        tmp.write('\\n'.join(unique_taxids) + '\\n')\n        tmp_name = tmp.name\n    \n    # -c shows the status code or the NEW TaxID if merged\n    cmd_update = f\"cat {tmp_name} | {cfg.TAXONKIT_BIN} lineage -c | cut -f 1,2\"\n    proc = subprocess.run(cmd_update, shell=True, capture_output=True, text=True)\n    \n    updated_ids = []\n    for line in proc.stdout.strip().split('\\n'):\n        if not line: continue\n        parts = line.split('\\t')\n        old, new = parts[0], parts[1]\n        # If new > 0, it's a valid or merged ID. If -1, it's not found.\n        final_id = new if int(new) > 0 else old\n        id_map[old] = final_id\n        updated_ids.append(final_id)\n        \n    # Using semicolons as a safer separator for shell/python splitting\n    tk_format = \"{domain|acellular root|superkingdom};{phylum};{class};{order};{family};{genus};{species}\"\n\n    for i in tqdm(range(0, len(unique_taxids), cfg.BATCH_SIZE), desc=\"Batch Query\"):\n        batch = unique_taxids[i : i + cfg.BATCH_SIZE]\n        with tempfile.NamedTemporaryFile(mode='w', delete=False) as tmp:\n            tmp.write('\\n'.join(batch) + '\\n')\n            tmp_name = tmp.name\n            \n        try:\n            # We use semicolon internally but TaxonKit outputs TAB between input and reformatted result\n            cmd = f\"cat {tmp_name} | {cfg.TAXONKIT_BIN} reformat2 --data-dir {cfg.DATA_DIR} -I 1 -f '{tk_format}' -r '0'\"\n            proc = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n            \n            for line in proc.stdout.strip().split('\\n'):\n                if not line: continue\n                parts = line.split('\\t')\n                if len(parts) < 2: continue\n                \n                tid = parts[0]\n                ranks = parts[1].split(';')\n                if len(ranks) < 7: continue \n                \n                results[tid] = {name: val for name, val in zip(cfg.RANK_NAMES, ranks)}\n        finally:\n            if os.path.exists(tmp_name): os.remove(tmp_name)\n\n    return results\n\n# ==================================================================================\n# 5. ENCODER\n# ==================================================================================\nclass TaxonomyEncoder:\n    def __init__(self, ranks):\n        self.ranks = ranks\n        self.vocab = {} \n        self.cards = []\n        \n    def fit(self, lineages_list, taxids_list):\n        # 1. Fit Ranks\n        for r in self.ranks:\n            values = set()\n            for l in lineages_list:\n                val = l.get(r, \"0\")\n                if val == \"\" or val == \"0\": val = cfg.UNKNOWN_TOKEN\n                values.add(val)\n            \n            unique = sorted(list(values - {cfg.UNKNOWN_TOKEN}))\n            mapping = {cfg.UNKNOWN_TOKEN: 0}\n            for idx, val in enumerate(unique):\n                mapping[val] = idx + 1\n            self.vocab[r] = mapping\n            self.cards.append(len(mapping))\n            \n        # 2. Fit Raw TaxID column\n        u_tids = sorted(list(set(taxids_list)))\n        t_map = {cfg.UNKNOWN_TOKEN: 0}\n        for i, t in enumerate(u_tids):\n            t_map[t] = i + 1\n        self.vocab['taxid'] = t_map\n        self.cards.append(len(t_map))\n            \n    def transform(self, lineages_list, taxids_list):\n        N = len(lineages_list)\n        out = np.zeros((N, len(self.ranks) + 1), dtype=np.int64)\n        \n        for i in range(N):\n            lin = lineages_list[i]\n            for j, r in enumerate(self.ranks):\n                val = lin.get(r, \"0\")\n                if val == \"\" or val == \"0\": val = cfg.UNKNOWN_TOKEN\n                out[i, j] = self.vocab[r].get(val, 0)\n            \n            tid = taxids_list[i]\n            out[i, -1] = self.vocab['taxid'].get(tid, 0)\n        return out\n\n# ==================================================================================\n# 6. EXECUTION\n# ==================================================================================\ndef main():\n    setup_taxonkit()\n    \n    print(\"\\n1️⃣  Parsing FASTAs...\")\n    train_map = extract_taxids_from_fasta(os.path.join(cfg.RAW_DATA_DIR, 'Train/train_sequences.fasta'))\n    test_map = extract_taxids_from_fasta(os.path.join(cfg.RAW_DATA_DIR, 'Test/testsuperset.fasta'))\n    \n    print(\"\\n2️⃣  Loading IDs...\")\n    train_ids = load_clean_ids(os.path.join(cfg.IDS_INPUT_DIR, 'train_ids.npy'))\n    test_ids = load_clean_ids(os.path.join(cfg.IDS_INPUT_DIR, 'test_ids.npy'))\n    \n    train_taxids = [train_map.get(pid, cfg.ROOT_TAXID) for pid in train_ids]\n    test_taxids = [test_map.get(pid, cfg.ROOT_TAXID) for pid in test_ids]\n    \n    print(\"\\n3️⃣  NCBI Lookup...\")\n    all_taxids = train_taxids + test_taxids\n    lineage_db = get_lineages_reformat2(all_taxids)\n    \n    train_lineages = [lineage_db.get(t, {}) for t in train_taxids]\n    test_lineages = [lineage_db.get(t, {}) for t in test_taxids]\n    \n    print(\"\\n4️⃣  Encoding...\")\n    encoder = TaxonomyEncoder(cfg.RANK_NAMES)\n    encoder.fit(train_lineages, train_taxids) # Fit on Train\n    \n    train_features = encoder.transform(train_lineages, train_taxids)\n    test_features = encoder.transform(test_lineages, test_taxids)\n    \n    # Final Output\n    np.save('train_tax_features.npy', train_features)\n    np.save('test_tax_features.npy', test_features)\n    np.save('tax_cardinalities.npy', np.array(encoder.cards))\n    \n    print(f\"✅ Finished. Train: {train_features.shape}, Test: {test_features.shape}\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-27T17:35:43.936432Z","iopub.execute_input":"2025-12-27T17:35:43.936703Z","iopub.status.idle":"2025-12-27T17:35:58.123570Z","shell.execute_reply.started":"2025-12-27T17:35:43.936668Z","shell.execute_reply":"2025-12-27T17:35:58.122658Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\n\n# Configuration (Paths from Kaggle environment)\nTAX_DIR = ''\nTRAIN_TAX_PATH = os.path.join(TAX_DIR, 'train_tax_features.npy')\nTEST_TAX_PATH = os.path.join(TAX_DIR, 'test_tax_features.npy')\n\ndef analyze_tax_col_7():\n    print(\"📂 Loading Taxonomy Features...\")\n    \n    # Check if files exist before loading\n    if not os.path.exists(TRAIN_TAX_PATH) or not os.path.exists(TEST_TAX_PATH):\n        print(f\"❌ Error: Could not find files in {TAX_DIR}\")\n        return\n\n    # 1. Load the raw numpy arrays\n    train_tax = np.load(TRAIN_TAX_PATH)\n    test_tax = np.load(TEST_TAX_PATH)\n\n    # 2. VIEW FIRST 15 ROWS (RAW DATA)\n    # We use a DataFrame for a cleaner visual representation of the rows/columns\n    print(\"\\n--- [PREVIEW] FIRST 15 ROWS OF RAW TRAIN DATA ---\")\n    raw_preview = pd.DataFrame(train_tax[:15])\n    print(raw_preview)\n    \n    # 3. Extract Column\n    COL_NUM = 1\n    train_types = train_tax[:, COL_NUM]\n    test_types = test_tax[:, COL_NUM]\n    \n    def get_stats(data, label):\n        unique, counts = np.unique(data, return_counts=True)\n        total = len(data)\n        \n        # Create statistics dataframe\n        df = pd.DataFrame({\n            'Tax_ID': unique.astype(int),\n            'Count': counts,\n            'Percentage': (counts / total) * 100\n        }).sort_values('Count', ascending=False)\n        \n        print(f\"\\n📊 {label} Distribution (Top 15 Most Frequent):\")\n        print(\"-\" * 55)\n        \n        # 4. SHOW TOP 40 ROWS OF STATS\n        top_15 = df.head(40)\n        for _, row in top_15.iterrows():\n            print(f\"Tax ID {int(row['Tax_ID']):<3} | {int(row['Count']):>8} proteins | {row['Percentage']:>6.2f}%\")\n        \n        return df\n\n    # Get stats for both sets\n    train_df = get_stats(train_types, \"TRAIN SET\")\n    test_df = get_stats(test_types, \"TEST SET\")\n    \n    # 5. Analysis & Skew Check\n    print(\"\\n\" + \"=\"*50)\n    print(\"🔍 FINAL ANALYSIS RESULTS\")\n    print(\"=\"*50)\n    \n    top_pct = train_df['Percentage'].iloc[0]\n    bot_pct = train_df['Percentage'].iloc[-1]\n    skew_ratio = top_pct / (bot_pct + 1e-6)\n    \n    print(f\"Total Unique Tax IDs found: {len(train_df)}\")\n    print(f\"Skew Ratio (Most Common vs Least Common): {skew_ratio:.2f}x\")\n    \n    if skew_ratio > 10:\n        print(\"\\n⚠️  STATUS: HIGH SKEW DETECTED\")\n        print(\"💡 RECOMMENDATION: Use the Lyapunov Stability Coef (V) >= 1.5 in your MoE Router.\")\n    else:\n        print(\"\\n✅ STATUS: BALANCED DATA\")\n        print(\"💡 RECOMMENDATION: Standard routing should work fine.\")\n\nif __name__ == \"__main__\":\n    # Ensure Pandas shows all columns if the array is wide\n    pd.set_option('display.max_columns', 20)\n    pd.set_option('display.width', 1000)\n    \n    analyze_tax_col_7()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-27T18:06:24.828660Z","iopub.execute_input":"2025-12-27T18:06:24.829034Z","iopub.status.idle":"2025-12-27T18:06:24.864404Z","shell.execute_reply.started":"2025-12-27T18:06:24.829008Z","shell.execute_reply":"2025-12-27T18:06:24.863501Z"}},"outputs":[],"execution_count":null}]}