{"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":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# CAFA 6 Protein Function Prediction - Data Understanding and EDA","metadata":{}},{"cell_type":"markdown","source":"## Competition Overview\nThis is a protein function prediction competition where we need to predict Gene Ontology (GO) terms for proteins based on their amino acid sequences. The competition uses a prospective evaluation approach - the test set will consist of proteins that gain experimental annotations between the submission deadline (February 2, 2026) and final evaluation (June 1, 2026). We need to predict GO terms across three subontologies: Molecular Function (MF), Biological Process (BP), and Cellular Component (CC).\n","metadata":{}},{"cell_type":"markdown","source":"## Understanding the Input Data\nThe model receives protein amino acid sequences as input. Each sequence is a string of characters representing the 20 standard amino acids. Our task is to predict which GO terms (functional annotations) apply to each protein, along with confidence scores between 0 and 1. The predictions must cover all three GO subontologies but will be evaluated separately for each.","metadata":{}},{"cell_type":"markdown","source":"## Understanding the Output Format\nThe submission format requires: ProteinID \\t GO_Term \\t Confidence_Score. Each protein can have multiple GO term predictions (up to 1500 combined across all three subontologies). Confidence scores must be in (0, 1.000] with up to 3 significant figures. If predictions aren't propagated to the ontology root, they'll be automatically propagated using maximum scores from children terms.","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom collections import Counter, defaultdict\nimport warnings\n!pip install Bio\nwarnings.filterwarnings('ignore')\n\nplt.style.use('seaborn-v0_8-darkgrid')\nsns.set_palette(\"husl\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:19:58.754082Z","iopub.execute_input":"2026-01-14T17:19:58.754819Z","iopub.status.idle":"2026-01-14T17:20:01.895391Z","shell.execute_reply.started":"2026-01-14T17:19:58.754786Z","shell.execute_reply":"2026-01-14T17:20:01.894542Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nThe training terms file contains our ground truth labels. Each row represents one GO term annotation for a protein, validated by experimental evidence. This is a multi-label classification problem where each protein can have multiple GO terms. The ontology column indicates which of the three subontologies (MFO, BPO, CCO) each term belongs to - this is crucial because evaluation happens separately for each subontology.","metadata":{}},{"cell_type":"code","source":"train_terms = pd.read_csv('/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv', \n                          sep='\\t', header=0, names=['protein_id', 'go_term', 'ontology'])\n\ntrain_terms = train_terms[train_terms['protein_id'] != 'EntryID']\n\nontology_mapping = {'P': 'BPO', 'C': 'CCO', 'F': 'MFO'}\ntrain_terms['ontology'] = train_terms['ontology'].map(ontology_mapping)\n\nprint(f\"Training annotations shape: {train_terms.shape}\")\nprint(f\"\\nFirst few annotations:\")\nprint(train_terms.head(10))\nprint(f\"\\nUnique proteins: {train_terms['protein_id'].nunique()}\")\nprint(f\"Unique GO terms: {train_terms['go_term'].nunique()}\")\nprint(f\"\\nOntology distribution:\")\nprint(train_terms['ontology'].value_counts())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:01.897340Z","iopub.execute_input":"2026-01-14T17:20:01.897968Z","iopub.status.idle":"2026-01-14T17:20:02.506587Z","shell.execute_reply.started":"2026-01-14T17:20:01.897935Z","shell.execute_reply":"2026-01-14T17:20:02.505648Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nUnderstanding the label distribution is critical for model design. Most proteins have multiple functions (multi-label), and the distribution shows significant variance. Some proteins have very few annotations while others have hundreds. This imbalance affects our loss function choice and sampling strategy - we need to handle proteins with varying annotation densities effectively during training.","metadata":{}},{"cell_type":"code","source":"annotations_per_protein = train_terms.groupby('protein_id').size()\n\nprint(f\"Annotation statistics per protein:\")\nprint(f\"Mean: {annotations_per_protein.mean():.2f}\")\nprint(f\"Median: {annotations_per_protein.median():.2f}\")\nprint(f\"Min: {annotations_per_protein.min()}\")\nprint(f\"Max: {annotations_per_protein.max()}\")\nprint(f\"Std: {annotations_per_protein.std():.2f}\")\n\nplt.figure(figsize=(14, 5))\n\nplt.subplot(1, 2, 1)\nplt.hist(annotations_per_protein, bins=50, edgecolor='black', alpha=0.7)\nplt.xlabel('Number of GO Terms per Protein')\nplt.ylabel('Frequency')\nplt.title('Distribution of Annotations per Protein')\nplt.axvline(annotations_per_protein.mean(), color='red', linestyle='--', label=f'Mean: {annotations_per_protein.mean():.1f}')\nplt.legend()\n\nplt.subplot(1, 2, 2)\nplt.hist(annotations_per_protein, bins=50, edgecolor='black', alpha=0.7, cumulative=True, density=True)\nplt.xlabel('Number of GO Terms per Protein')\nplt.ylabel('Cumulative Probability')\nplt.title('Cumulative Distribution of Annotations')\nplt.grid(True, alpha=0.3)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:02.507683Z","iopub.execute_input":"2026-01-14T17:20:02.508035Z","iopub.status.idle":"2026-01-14T17:20:03.207717Z","shell.execute_reply.started":"2026-01-14T17:20:02.508010Z","shell.execute_reply":"2026-01-14T17:20:03.206717Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nEach subontology (BPO, CCO, MFO) has different annotation patterns. Biological Process terms tend to be more numerous per protein than Cellular Component terms. This matters because we'll be evaluated separately on each subontology - we might need different prediction thresholds or strategies for each one to maximize the F1-score independently.","metadata":{}},{"cell_type":"code","source":"ontology_annotations = train_terms.groupby(['protein_id', 'ontology']).size().unstack(fill_value=0)\n\nprint(\"Annotations per protein by ontology:\")\nprint(ontology_annotations.describe())\n\nfig, axes = plt.subplots(1, 3, figsize=(15, 5))\n\nfor idx, ont in enumerate(['BPO', 'CCO', 'MFO']):\n    if ont in ontology_annotations.columns:\n        axes[idx].hist(ontology_annotations[ont], bins=30, edgecolor='black', alpha=0.7)\n        axes[idx].set_xlabel(f'Number of {ont} Terms')\n        axes[idx].set_ylabel('Frequency')\n        axes[idx].set_title(f'{ont} Annotations Distribution')\n        axes[idx].axvline(ontology_annotations[ont].mean(), color='red', linestyle='--', \n                         label=f'Mean: {ontology_annotations[ont].mean():.1f}')\n        axes[idx].legend()\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:03.210100Z","iopub.execute_input":"2026-01-14T17:20:03.210401Z","iopub.status.idle":"2026-01-14T17:20:04.233781Z","shell.execute_reply.started":"2026-01-14T17:20:03.210376Z","shell.execute_reply":"2026-01-14T17:20:04.232582Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nGO term frequency follows a heavy-tailed distribution - a few terms are extremely common while many are rare. This extreme class imbalance is crucial for model training. Rare terms are harder to predict but have higher information content (weights) in evaluation. We need techniques like focal loss, class weighting, or specialized sampling to ensure our model learns rare terms effectively, not just the common ones.","metadata":{}},{"cell_type":"code","source":"term_frequency = train_terms['go_term'].value_counts()\n\nprint(f\"GO term frequency statistics:\")\nprint(f\"Total unique terms: {len(term_frequency)}\")\nprint(f\"Terms appearing once: {(term_frequency == 1).sum()}\")\nprint(f\"Terms appearing >100 times: {(term_frequency > 100).sum()}\")\nprint(f\"Terms appearing >1000 times: {(term_frequency > 1000).sum()}\")\n\nprint(f\"\\nTop 15 most common GO terms:\")\nprint(term_frequency.head(15))\n\nplt.figure(figsize=(14, 5))\n\nplt.subplot(1, 2, 1)\nplt.hist(np.log10(term_frequency + 1), bins=50, edgecolor='black', alpha=0.7)\nplt.xlabel('Log10(Frequency + 1)')\nplt.ylabel('Number of GO Terms')\nplt.title('GO Term Frequency Distribution (log scale)')\n\nplt.subplot(1, 2, 2)\ncumsum = np.cumsum(term_frequency.values)\nplt.plot(cumsum / cumsum[-1])\nplt.xlabel('GO Term Rank')\nplt.ylabel('Cumulative Proportion of Annotations')\nplt.title('Cumulative Distribution of Annotations by Term')\nplt.grid(True, alpha=0.3)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:04.235188Z","iopub.execute_input":"2026-01-14T17:20:04.235573Z","iopub.status.idle":"2026-01-14T17:20:04.686873Z","shell.execute_reply.started":"2026-01-14T17:20:04.235539Z","shell.execute_reply":"2026-01-14T17:20:04.686002Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nProtein sequences vary dramatically in length from very short peptides to extremely long proteins. This variation is critical for model architecture decisions. If using transformer models, we need to handle variable-length inputs efficiently - either through padding/truncation, sliding windows for long sequences, or hierarchical approaches. Length distribution also affects batch composition and memory requirements.","metadata":{}},{"cell_type":"code","source":"from Bio import SeqIO\n\ntrain_sequences = {}\nsequence_lengths = []\n\nwith open('/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta', 'r') as f:\n    for record in SeqIO.parse(f, 'fasta'):\n        protein_id = record.id.split('|')[1]\n        sequence = str(record.seq)\n        train_sequences[protein_id] = sequence\n        sequence_lengths.append(len(sequence))\n\nprint(f\"Total training sequences: {len(train_sequences)}\")\nprint(f\"\\nSequence length statistics:\")\nprint(f\"Mean: {np.mean(sequence_lengths):.2f}\")\nprint(f\"Median: {np.median(sequence_lengths):.2f}\")\nprint(f\"Min: {np.min(sequence_lengths)}\")\nprint(f\"Max: {np.max(sequence_lengths)}\")\nprint(f\"Std: {np.std(sequence_lengths):.2f}\")\n\nplt.figure(figsize=(14, 5))\n\nplt.subplot(1, 2, 1)\nplt.hist(sequence_lengths, bins=50, edgecolor='black', alpha=0.7)\nplt.xlabel('Sequence Length')\nplt.ylabel('Frequency')\nplt.title('Distribution of Protein Sequence Lengths')\nplt.axvline(np.mean(sequence_lengths), color='red', linestyle='--', label=f'Mean: {np.mean(sequence_lengths):.1f}')\nplt.legend()\n\nplt.subplot(1, 2, 2)\nplt.hist(sequence_lengths, bins=50, edgecolor='black', alpha=0.7, cumulative=True, density=True)\nplt.xlabel('Sequence Length')\nplt.ylabel('Cumulative Probability')\nplt.title('Cumulative Distribution of Sequence Lengths')\nplt.grid(True, alpha=0.3)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:04.688018Z","iopub.execute_input":"2026-01-14T17:20:04.688312Z","iopub.status.idle":"2026-01-14T17:20:14.896877Z","shell.execute_reply.started":"2026-01-14T17:20:04.688288Z","shell.execute_reply":"2026-01-14T17:20:14.895805Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nAmino acid composition shows which building blocks are most common in proteins. Some amino acids like Leucine and Alanine are very frequent while others like Tryptophan are rare. This informs our tokenization strategy - we might use character-level encoding (21 tokens: 20 amino acids + special tokens) or k-mer encoding. Understanding base frequencies helps with model initialization and rare token handling.","metadata":{}},{"cell_type":"code","source":"amino_acid_composition = Counter()\namino_acids = set()\n\nfor seq in train_sequences.values():\n    amino_acid_composition.update(seq)\n    amino_acids.update(seq)\n\ntotal_aa = sum(amino_acid_composition.values())\naa_freq = {aa: count/total_aa for aa, count in amino_acid_composition.items()}\n\nprint(f\"Unique amino acids found: {len(amino_acids)}\")\nprint(f\"Standard 20 amino acids: {amino_acids}\")\nprint(f\"\\nAmino acid frequency (%):\")\nfor aa, freq in sorted(aa_freq.items(), key=lambda x: x[1], reverse=True):\n    print(f\"{aa}: {freq*100:.2f}%\")\n\nplt.figure(figsize=(12, 6))\nsorted_aa = sorted(aa_freq.items(), key=lambda x: x[1], reverse=True)\naas, freqs = zip(*sorted_aa)\nplt.bar(aas, [f*100 for f in freqs], edgecolor='black', alpha=0.7)\nplt.xlabel('Amino Acid')\nplt.ylabel('Frequency (%)')\nplt.title('Amino Acid Composition in Training Set')\nplt.grid(True, alpha=0.3, axis='y')\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:14.898399Z","iopub.execute_input":"2026-01-14T17:20:14.898731Z","iopub.status.idle":"2026-01-14T17:20:18.845702Z","shell.execute_reply.started":"2026-01-14T17:20:14.898695Z","shell.execute_reply":"2026-01-14T17:20:18.844635Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nExamining individual examples shows the actual input-output relationship. Each protein sequence (a string of amino acids) maps to multiple GO terms across different ontologies. This is what our model must learn - the relationship between sequence patterns and functional annotations. The multi-label, multi-ontology nature means we need a model architecture that can output predictions for thousands of possible GO terms simultaneously.","metadata":{}},{"cell_type":"code","source":"print(\"Checking ontology column values:\")\nprint(train_terms['ontology'].unique())\nprint(f\"\\nValue counts:\")\nprint(train_terms['ontology'].value_counts())\n\nsample_proteins = list(train_sequences.keys())[:5]\n\nprint(\"\\n\" + \"=\"*80)\nprint(\"Sample protein sequences and their annotations:\")\nprint(\"=\"*80 + \"\\n\")\n\nfor protein_id in sample_proteins:\n    sequence = train_sequences[protein_id]\n    annotations = train_terms[train_terms['protein_id'] == protein_id]\n    \n    print(f\"Protein: {protein_id}\")\n    print(f\"Sequence length: {len(sequence)}\")\n    print(f\"Sequence (first 60 chars): {sequence[:60]}...\")\n    print(f\"Total annotations: {len(annotations)}\")\n    \n    if len(annotations) > 0:\n        print(f\"Ontologies present: {annotations['ontology'].unique()}\")\n        print(f\"GO terms: {annotations['go_term'].tolist()[:5]}\")\n        \n        for ont in ['BPO', 'CCO', 'MFO']:\n            ont_terms = annotations[annotations['ontology'] == ont]['go_term'].tolist()\n            if ont_terms:\n                print(f\"  {ont}: {len(ont_terms)} terms - {ont_terms[:3]}\")\n    else:\n        print(\"  No annotations found\")\n    print()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:18.846906Z","iopub.execute_input":"2026-01-14T17:20:18.847275Z","iopub.status.idle":"2026-01-14T17:20:19.175572Z","shell.execute_reply.started":"2026-01-14T17:20:18.847249Z","shell.execute_reply":"2026-01-14T17:20:19.174391Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nTaxonomic information reveals species distribution in the training data. Some species (like human, mouse, yeast) are heavily represented while others are sparse. This creates a domain shift challenge - proteins from well-studied species might have different annotation patterns than those from rare species. We could use taxon ID as an auxiliary input feature or apply domain adaptation techniques to improve generalization across species.","metadata":{}},{"cell_type":"code","source":"train_taxonomy = pd.read_csv('/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv', sep='\\t', header=None, \n                             names=['protein_id', 'taxon_id'])\n\nprint(f\"Taxonomy information shape: {train_taxonomy.shape}\")\nprint(f\"Unique species (taxon IDs): {train_taxonomy['taxon_id'].nunique()}\")\nprint(f\"\\nTop 10 species by protein count:\")\nprint(train_taxonomy['taxon_id'].value_counts().head(10))\n\nspecies_distribution = train_taxonomy['taxon_id'].value_counts()\n\nplt.figure(figsize=(12, 6))\nplt.bar(range(len(species_distribution.head(20))), species_distribution.head(20).values, \n        edgecolor='black', alpha=0.7)\nplt.xlabel('Species Rank')\nplt.ylabel('Number of Proteins')\nplt.title('Top 20 Species by Protein Count')\nplt.xticks(range(len(species_distribution.head(20))), \n           [f'Taxon {tid}' for tid in species_distribution.head(20).index], rotation=45, ha='right')\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:19.177435Z","iopub.execute_input":"2026-01-14T17:20:19.177759Z","iopub.status.idle":"2026-01-14T17:20:19.525686Z","shell.execute_reply.started":"2026-01-14T17:20:19.177733Z","shell.execute_reply":"2026-01-14T17:20:19.524650Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nInformation Accretion weights determine how much each GO term contributes to the evaluation metric. Rare, specific terms have higher weights than common, generic terms. This is critical for model optimization - we should incorporate these weights into our loss function or post-processing. Simply maximizing unweighted F1 won't give us the best competition score; we need to prioritize predictions weighted by IA values.","metadata":{}},{"cell_type":"code","source":"ia_weights = pd.read_csv('/kaggle/input/cafa-6-protein-function-prediction/IA.tsv', sep='\\t', header=None, \n                         names=['go_term', 'ia_weight'])\n\nprint(f\"Information Accretion weights shape: {ia_weights.shape}\")\nprint(f\"\\nIA weight statistics:\")\nprint(ia_weights['ia_weight'].describe())\n\ncommon_terms = set(term_frequency.index) & set(ia_weights['go_term'])\nprint(f\"\\nGO terms with both frequency and IA weight data: {len(common_terms)}\")\nprint(f\"Terms in training without IA weights: {len(set(term_frequency.index) - set(ia_weights['go_term']))}\")\n\nia_dict = ia_weights.set_index('go_term')['ia_weight'].to_dict()\nmatched_frequencies = []\nmatched_ia_weights = []\n\nfor term in term_frequency.index:\n    if term in ia_dict:\n        matched_frequencies.append(term_frequency[term])\n        matched_ia_weights.append(ia_dict[term])\n\nplt.figure(figsize=(14, 5))\n\nplt.subplot(1, 2, 1)\nplt.hist(ia_weights['ia_weight'], bins=50, edgecolor='black', alpha=0.7)\nplt.xlabel('Information Accretion Weight')\nplt.ylabel('Frequency')\nplt.title('Distribution of IA Weights')\nplt.axvline(ia_weights['ia_weight'].mean(), color='red', linestyle='--', \n            label=f'Mean: {ia_weights[\"ia_weight\"].mean():.2f}')\nplt.legend()\n\nplt.subplot(1, 2, 2)\nplt.scatter(matched_frequencies, matched_ia_weights, alpha=0.3)\nplt.xlabel('Term Frequency (log scale)')\nplt.ylabel('IA Weight')\nplt.xscale('log')\nplt.title('Relationship between Term Frequency and IA Weight')\nplt.grid(True, alpha=0.3)\n\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:19.528213Z","iopub.execute_input":"2026-01-14T17:20:19.528498Z","iopub.status.idle":"2026-01-14T17:20:20.760433Z","shell.execute_reply.started":"2026-01-14T17:20:19.528475Z","shell.execute_reply":"2026-01-14T17:20:20.759502Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nThe test superset contains all proteins we need to make predictions for, though only a subset will actually be evaluated (those gaining experimental annotations during the curation period). The sequence length distribution should be similar to training data. Any significant distribution shift could indicate domain adaptation challenges. We must generate predictions for every protein in this file, even though most won't be in the final test set.","metadata":{}},{"cell_type":"code","source":"test_superset_sequences = {}\ntest_taxon_ids = {}\n\nwith open('/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta', 'r') as f:\n    for record in SeqIO.parse(f, 'fasta'):\n        header_parts = record.description.split()\n        protein_id = record.id\n        taxon_id = None\n        \n        for part in header_parts:\n            if part.startswith('OX='):\n                taxon_id = part.split('=')[1]\n                break\n        \n        test_superset_sequences[protein_id] = str(record.seq)\n        test_taxon_ids[protein_id] = taxon_id\n\nprint(f\"Test superset size: {len(test_superset_sequences)}\")\nprint(f\"\\nTest sequence length statistics:\")\ntest_lengths = [len(seq) for seq in test_superset_sequences.values()]\nprint(f\"Mean: {np.mean(test_lengths):.2f}\")\nprint(f\"Median: {np.median(test_lengths):.2f}\")\nprint(f\"Min: {np.min(test_lengths)}\")\nprint(f\"Max: {np.max(test_lengths)}\")\n\nplt.figure(figsize=(12, 6))\nplt.hist(test_lengths, bins=50, edgecolor='black', alpha=0.7, label='Test')\nplt.hist(sequence_lengths, bins=50, edgecolor='black', alpha=0.5, label='Train')\nplt.xlabel('Sequence Length')\nplt.ylabel('Frequency')\nplt.title('Comparison of Sequence Length Distributions')\nplt.legend()\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:20.761514Z","iopub.execute_input":"2026-01-14T17:20:20.761856Z","iopub.status.idle":"2026-01-14T17:20:24.324478Z","shell.execute_reply.started":"2026-01-14T17:20:20.761829Z","shell.execute_reply":"2026-01-14T17:20:24.323510Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nComparing train and test species distributions reveals whether we face a zero-shot species generalization problem. If test proteins come from species not in training, our model must generalize across evolutionary distances. This strongly suggests using protein language models (ESM, ProtBERT) pretrained on diverse protein sequences rather than training from scratch on this limited dataset.","metadata":{}},{"cell_type":"code","source":"test_taxon = pd.read_csv('/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv', \n                         sep='\\t', header=0)\n\nprint(f\"Unique species in test superset: {test_taxon['ID'].nunique()}\")\nprint(f\"\\nTop 10 test species:\")\nprint(test_taxon.head(10))\n\ntrain_species = set(train_taxonomy['taxon_id'].unique())\ntest_species = set(test_taxon['ID'].unique())\n\nprint(f\"\\nSpecies overlap:\")\nprint(f\"Species only in train: {len(train_species - test_species)}\")\nprint(f\"Species only in test: {len(test_species - train_species)}\")\nprint(f\"Species in both: {len(train_species & test_species)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:24.325717Z","iopub.execute_input":"2026-01-14T17:20:24.326053Z","iopub.status.idle":"2026-01-14T17:20:24.368020Z","shell.execute_reply.started":"2026-01-14T17:20:24.326026Z","shell.execute_reply":"2026-01-14T17:20:24.367119Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nThis reveals our output space size - thousands of possible GO terms across three ontologies. Our model architecture must output probability scores for all these terms. We could use three separate output heads (one per ontology) or a single unified head. The extreme number of output classes suggests using efficient multi-label classification techniques like binary cross-entropy with sigmoid activation rather than softmax.","metadata":{}},{"cell_type":"code","source":"term_ontology_map = train_terms.set_index('go_term')['ontology'].to_dict()\n\ntrain_terms_by_ontology = {\n    'BPO': set(train_terms[train_terms['ontology'] == 'BPO']['go_term'].unique()),\n    'CCO': set(train_terms[train_terms['ontology'] == 'CCO']['go_term'].unique()),\n    'MFO': set(train_terms[train_terms['ontology'] == 'MFO']['go_term'].unique())\n}\n\nprint(\"Unique GO terms by ontology in training:\")\nfor ont, terms in train_terms_by_ontology.items():\n    print(f\"{ont}: {len(terms)} unique terms\")\n\ntotal_possible_labels = sum(len(terms) for terms in train_terms_by_ontology.values())\nprint(f\"\\nTotal possible output labels: {total_possible_labels}\")\n\nplt.figure(figsize=(10, 6))\nontology_counts = [len(terms) for terms in train_terms_by_ontology.values()]\nplt.bar(train_terms_by_ontology.keys(), ontology_counts, edgecolor='black', alpha=0.7)\nplt.xlabel('Ontology')\nplt.ylabel('Number of Unique GO Terms')\nplt.title('GO Terms Distribution by Ontology')\nplt.grid(True, alpha=0.3, axis='y')\nplt.tight_layout()\nplt.show()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:24.369185Z","iopub.execute_input":"2026-01-14T17:20:24.369551Z","iopub.status.idle":"2026-01-14T17:20:25.254741Z","shell.execute_reply.started":"2026-01-14T17:20:24.369516Z","shell.execute_reply":"2026-01-14T17:20:25.254048Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nMerging all data sources creates our complete training dataset. Each row represents one protein-GO term pair with its ontology, species, and evaluation weight. This structure is perfect for creating training batches where each protein sequence maps to a multi-hot encoded label vector weighted by IA values. Missing sequences would be problematic and need investigation.","metadata":{}},{"cell_type":"code","source":"merged_train = train_terms.merge(train_taxonomy, on='protein_id')\nmerged_train = merged_train.merge(ia_weights, on='go_term')\n\nprint(\"Training data with all features merged:\")\nprint(merged_train.head())\nprint(f\"\\nShape: {merged_train.shape}\")\n\ncoverage_check = set(train_terms['protein_id'].unique()) - set(train_sequences.keys())\nprint(f\"\\nProteins in annotations but missing sequences: {len(coverage_check)}\")\n\nif len(coverage_check) > 0:\n    print(f\"Examples: {list(coverage_check)[:5]}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:25.255710Z","iopub.execute_input":"2026-01-14T17:20:25.256019Z","iopub.status.idle":"2026-01-14T17:20:25.667183Z","shell.execute_reply.started":"2026-01-14T17:20:25.255994Z","shell.execute_reply":"2026-01-14T17:20:25.666185Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<br><br>\nThis comprehensive specification documents exactly what our model needs to consume and produce. The model takes variable-length amino acid sequences and outputs probability distributions over thousands of GO terms. The key challenge is multi-label classification with extreme class imbalance, requiring careful loss design and threshold optimization. Using pretrained protein language models is almost certainly necessary given the limited training data and need for cross-species generalization.","metadata":{}},{"cell_type":"code","source":"print(\"=\" * 80)\nprint(\"MODEL INPUT-OUTPUT SPECIFICATION\")\nprint(\"=\" * 80)\n\nprint(\"\\nMODEL INPUTS:\")\nprint(\"-\" * 80)\nprint(\"1. Protein Amino Acid Sequence (string)\")\nprint(f\"   - Variable length: {np.min(sequence_lengths)} to {np.max(sequence_lengths)} amino acids\")\nprint(f\"   - Alphabet size: 20 standard amino acids + special tokens (found {len(amino_acids)} unique)\")\nprint(\"   - Example format: 'MKTAYIAKQRQISFVKSHFSRQLEERLGLIEVQAPILSRVGD...'\")\n\nprint(\"\\n2. Optional: Taxonomic ID (integer)\")\nprint(f\"   - Range: Various species identifiers\")\nprint(\"   - Can be used as auxiliary feature or for domain adaptation\")\n\nprint(\"\\n\\nMODEL OUTPUTS:\")\nprint(\"-\" * 80)\nprint(f\"Multi-label binary classification for {total_possible_labels} GO terms\")\n\nfor ont, terms in train_terms_by_ontology.items():\n    print(f\"\\n{ont} ({len(terms)} terms):\")\n    print(f\"   - Output: Probability scores in (0, 1] for each term\")\n    print(f\"   - Format: Sigmoid activation (not softmax)\")\n    print(f\"   - Evaluation: Weighted F1 using IA weights\")\n\nprint(\"\\n\\nSUBMISSION FORMAT:\")\nprint(\"-\" * 80)\nprint(\"Format: ProteinID <TAB> GO_Term <TAB> Confidence_Score\")\nprint(\"Requirements:\")\nprint(\"  - Confidence in (0, 1.000] with up to 3 significant figures\")\nprint(\"  - Maximum 1500 predictions per protein (combined across ontologies)\")\nprint(\"  - Only include predictions > 0 (don't predict 0)\")\nprint(\"  - Predictions auto-propagated to ontology root if needed\")\n\nprint(\"\\n\\nTRAINING STRATEGY:\")\nprint(\"-\" * 80)\nprint(\"1. Use pretrained protein language models (ESM-2, ProtBERT)\")\nprint(\"2. Fine-tune with multi-label binary cross-entropy\")\nprint(\"3. Apply class weights based on IA values\")\nprint(\"4. Handle class imbalance (focal loss, oversampling rare terms)\")\nprint(\"5. Separate heads or strategies for each ontology\")\nprint(\"6. Optimize threshold per ontology to maximize weighted F1\")\n\nprint(\"\\n\" + \"=\" * 80)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-14T17:20:25.668616Z","iopub.execute_input":"2026-01-14T17:20:25.669483Z","iopub.status.idle":"2026-01-14T17:20:25.690822Z","shell.execute_reply.started":"2026-01-14T17:20:25.669441Z","shell.execute_reply":"2026-01-14T17:20:25.689703Z"}},"outputs":[],"execution_count":null}]}