{"cells": [{"cell_type": "markdown", "metadata": {}, "source": "# CAFA-6 Hybrid: Diamond BLASTP + Proven Ensemble\n\n**Strategy:**\n1. Use Diamond BLASTP for proteins with high-confidence homology hits\n2. Fall back to proven ensemble for proteins without Diamond hits\n3. Combine predictions with weighting"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "import os\nimport gc\nimport pandas as pd\nimport numpy as np\nfrom collections import defaultdict\nfrom tqdm.auto import tqdm\n\nprint('Hybrid Diamond+Ensemble Solution v1')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 1. SETUP DIAMOND\n# ==========================================\nprint('[1/9] Setting up Diamond...')\n\nif not os.path.exists('diamond'):\n    !wget -q https://github.com/bbuchfink/diamond/releases/download/v2.1.8/diamond-linux64.tar.gz\n    !tar xzf diamond-linux64.tar.gz\n    !chmod +x diamond\n    print('Diamond installed!')\nelse:\n    print('Diamond already exists')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 2. PATHS AND KERNEL SOURCES\n# ==========================================\nTRAIN_FASTA = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\nTEST_FASTA = '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta'\nTRAIN_TERMS = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv'\nGO_OBO = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo'\n\n# Proven ensemble kernels (from kernel sources)\nQUICK_KERNEL = '/kaggle/input/cafa-quick-submission-0-358-score/submission.tsv'\nPROVEN_KERNEL = '/kaggle/input/cafa-6-proven-ensemble-submission/submission.tsv'\n\n# Check files\nfor f in [TRAIN_FASTA, TEST_FASTA, TRAIN_TERMS, GO_OBO, QUICK_KERNEL, PROVEN_KERNEL]:\n    if os.path.exists(f):\n        print(f'OK: {os.path.basename(f)}')\n    else:\n        print(f'MISSING: {f}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 3. CREATE DIAMOND DATABASE\n# ==========================================\nprint('[2/9] Creating Diamond database...')\n\n!./diamond makedb --in {TRAIN_FASTA} -d train_db --quiet\n\nif os.path.exists('train_db.dmnd'):\n    print('Database created!')\nelse:\n    print('ERROR: Database creation failed')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 4. RUN DIAMOND BLASTP\n# ==========================================\nprint('[3/9] Running Diamond BLASTP...')\nprint('This may take 10-15 minutes...')\n\n!./diamond blastp \\\n    -d train_db \\\n    -q {TEST_FASTA} \\\n    --outfmt 6 qseqid sseqid pident evalue bitscore \\\n    --top 10 \\\n    --more-sensitive \\\n    --quiet \\\n    --out diamond_results.tsv\n\nif os.path.exists('diamond_results.tsv'):\n    results = pd.read_csv('diamond_results.tsv', sep='\\t', header=None,\n                          names=['query', 'target', 'pident', 'evalue', 'bitscore'])\n    print(f'Alignment complete! {len(results):,} hits')\n    print(f'Unique test proteins with hits: {results[\"query\"].nunique():,}')\nelse:\n    print('ERROR: Alignment failed')\n    results = pd.DataFrame()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 5. LOAD TRAINING GO TERMS\n# ==========================================\nprint('[4/9] Loading training GO terms...')\n\ntrain_terms = pd.read_csv(TRAIN_TERMS, sep='\\t', header=None, \n                          names=['protein', 'go_term', 'aspect'])\nprint(f'Loaded {len(train_terms):,} GO annotations')\n\n# Create protein -> GO terms mapping\nprotein_to_go = defaultdict(set)\nfor _, row in tqdm(train_terms.iterrows(), total=len(train_terms), desc='Building GO map'):\n    protein_to_go[row['protein']].add(row['go_term'])\n\nprint(f'GO map built for {len(protein_to_go):,} proteins')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 6. TRANSFER GO TERMS FROM DIAMOND HITS\n# ==========================================\nprint('[5/9] Transferring GO terms from Diamond hits...')\n\n# Filter by minimum identity\nMIN_IDENTITY = 40.0  # Higher threshold for more confidence\nresults_filtered = results[results['pident'] >= MIN_IDENTITY]\nprint(f'Hits after filtering (>={MIN_IDENTITY}% identity): {len(results_filtered):,}')\n\n# Transfer GO terms\ndiamond_predictions = defaultdict(lambda: defaultdict(float))\n\nfor _, row in tqdm(results_filtered.iterrows(), total=len(results_filtered), desc='Transferring'):\n    query = row['query']\n    target = row['target']\n    \n    # Parse UniProt ID from FASTA header format (sp|ID|NAME or tr|ID|NAME)\n    if '|' in target:\n        target = target.split('|')[1]\n    \n    pident = row['pident'] / 100.0\n    \n    if target in protein_to_go:\n        for go_term in protein_to_go[target]:\n            diamond_predictions[query][go_term] = max(diamond_predictions[query][go_term], pident)\n\nprint(f'Diamond predictions for {len(diamond_predictions):,} proteins')\ndiamond_proteins = set(diamond_predictions.keys())"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 7. LOAD PROVEN ENSEMBLE PREDICTIONS\n# ==========================================\nprint('[6/9] Loading proven ensemble predictions...')\n\ndef load_predictions(path):\n    df = pd.read_csv(path, sep='\\t', header=None, names=['protein', 'go_term', 'score'])\n    preds = defaultdict(lambda: defaultdict(float))\n    for _, row in tqdm(df.iterrows(), total=len(df), desc=f'Loading {os.path.basename(path)}'):\n        preds[row['protein']][row['go_term']] = max(preds[row['protein']][row['go_term']], row['score'])\n    return preds\n\nquick_preds = load_predictions(QUICK_KERNEL)\nproven_preds = load_predictions(PROVEN_KERNEL)\n\nprint(f'Quick predictions: {len(quick_preds):,} proteins')\nprint(f'Proven predictions: {len(proven_preds):,} proteins')\n\n# Create ensemble from quick + proven\nensemble_preds = defaultdict(lambda: defaultdict(float))\nall_ensemble_proteins = set(quick_preds.keys()) | set(proven_preds.keys())\n\nfor protein in tqdm(all_ensemble_proteins, desc='Merging ensemble'):\n    all_terms = set(quick_preds[protein].keys()) | set(proven_preds[protein].keys())\n    for term in all_terms:\n        s1 = quick_preds[protein].get(term, 0)\n        s2 = proven_preds[protein].get(term, 0)\n        # Simple average\n        ensemble_preds[protein][term] = (s1 + s2) / 2\n\nprint(f'Ensemble has {len(ensemble_preds):,} proteins')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 8. COMBINE DIAMOND + ENSEMBLE\n# ==========================================\nprint('[7/9] Combining Diamond + Ensemble...')\n\n# For proteins with Diamond hits: use Diamond with ensemble boost\n# For proteins without Diamond hits: use ensemble only\n\nfinal_predictions = defaultdict(lambda: defaultdict(float))\nDIAMOND_WEIGHT = 0.7  # Weight for Diamond predictions\nENSEMBLE_WEIGHT = 0.3  # Weight for ensemble predictions\n\n# Process proteins with Diamond hits\nfor protein in tqdm(diamond_proteins, desc='Diamond proteins'):\n    all_terms = set(diamond_predictions[protein].keys()) | set(ensemble_preds[protein].keys())\n    for term in all_terms:\n        d_score = diamond_predictions[protein].get(term, 0)\n        e_score = ensemble_preds[protein].get(term, 0)\n        # Weighted combination\n        final_predictions[protein][term] = DIAMOND_WEIGHT * d_score + ENSEMBLE_WEIGHT * e_score\n\n# Process proteins without Diamond hits (use ensemble)\nnon_diamond_proteins = set(ensemble_preds.keys()) - diamond_proteins\nfor protein in tqdm(non_diamond_proteins, desc='Ensemble-only proteins'):\n    for term, score in ensemble_preds[protein].items():\n        final_predictions[protein][term] = score\n\nprint(f'Combined predictions for {len(final_predictions):,} proteins')\nprint(f'  - Diamond+Ensemble: {len(diamond_proteins):,}')\nprint(f'  - Ensemble-only: {len(non_diamond_proteins):,}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 9. GO PROPAGATION\n# ==========================================\nprint('[8/9] GO ontology propagation...')\n\ndef parse_obo(obo_path):\n    parents = defaultdict(set)\n    with open(obo_path, 'r') as f:\n        cur_id = None\n        for line in f:\n            line = line.strip()\n            if line == '[Term]':\n                cur_id = None\n            elif line.startswith('id: '):\n                cur_id = line.split('id: ')[1].strip()\n            elif line.startswith('is_a: '):\n                pid = line.split()[1].strip()\n                if cur_id:\n                    parents[cur_id].add(pid)\n            elif line.startswith('relationship: part_of '):\n                parts = line.split()\n                if len(parts) >= 3:\n                    pid = parts[2].strip()\n                    if cur_id:\n                        parents[cur_id].add(pid)\n    return parents\n\ndef get_ancestors(term, parents_map, cache):\n    if term in cache:\n        return cache[term]\n    ancestors = set()\n    stack = [term]\n    while stack:\n        cur = stack.pop()\n        for p in parents_map.get(cur, []):\n            if p not in ancestors:\n                ancestors.add(p)\n                stack.append(p)\n    cache[term] = ancestors\n    return ancestors\n\nparents_map = parse_obo(GO_OBO)\nprint(f'Parsed {len(parents_map)} GO terms')\n\nROOTS = {'GO:0003674', 'GO:0008150', 'GO:0005575'}\nancestor_cache = {}\nfinal_rows = []\n\nfor protein, terms_dict in tqdm(final_predictions.items(), desc='Propagating'):\n    final_scores = dict(terms_dict)\n    \n    for term, score in list(terms_dict.items()):\n        for ancestor in get_ancestors(term, parents_map, ancestor_cache):\n            final_scores[ancestor] = max(final_scores.get(ancestor, 0), score)\n    \n    for root in ROOTS:\n        final_scores[root] = 1.0\n    \n    for term, score in final_scores.items():\n        if score >= 0.01:\n            final_rows.append((protein, term, round(score, 6)))\n\nprint(f'Total predictions after propagation: {len(final_rows):,}')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# ==========================================\n# 10. SAVE SUBMISSION\n# ==========================================\nprint('[9/9] Saving submission...')\n\nfinal_df = pd.DataFrame(final_rows, columns=['protein_id', 'go_term', 'score'])\nfinal_df = final_df.sort_values(['protein_id', 'score'], ascending=[True, False])\nfinal_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)\n\nprint(f'\\n{\"=\"*60}')\nprint('DONE - Hybrid Diamond+Ensemble Solution v1')\nprint(f'Predictions: {len(final_df):,}')\nprint(f'Proteins: {final_df[\"protein_id\"].nunique():,}')\nprint(f'{\"=\"*60}')\nprint(final_df.head(10))"}], "metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"}, "language_info": {"name": "python", "version": "3.11.0"}}, "nbformat": 4, "nbformat_minor": 4}