\n\n### The Challenge\n\n**CAFA (Critical Assessment of protein Function Annotation)** is a community challenge to evaluate and improve computational methods for predicting protein function. This is CAFA 6!\n\n**Our Task:** Given a protein sequence, predict its functional annotations using Gene Ontology (GO) terms across three aspects:\n\n| Aspect | Code | Description |\n|--------|------|-------------|\n| **Molecular Function** | MF | The biochemical activity of the protein |\n| **Biological Process** | BP | The larger biological goal the protein contributes to |\n| **Cellular Component** | CC | The location in the cell where the protein acts |\n\n### Evaluation Metric\n\nThe competition uses a **weighted F-max score** that accounts for:\n1. Precision and recall across different confidence thresholds\n2. **Information Accretion (IA)** - weights based on term specificity\n3. GO hierarchy constraints (True Path Rule)\n\n$$F_{max} = \\max_\\tau \\frac{2 \\cdot pr(\\tau) \\cdot rc(\\tau)}{pr(\\tau) + rc(\\tau)}$$\n\n
"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 2. Setup & Configuration"}, {"cell_type": "code", "source": "# Install required packages (Kaggle environment)\n!pip install -q biopython goatools", "metadata": {}, "execution_count": null, "outputs": []}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Core libraries\nimport numpy as np\nimport pandas as pd\nimport os\nimport gc\nimport warnings\nfrom collections import Counter, defaultdict\nfrom tqdm.auto import tqdm\n\n# Visualization\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nimport plotly.express as px\nimport plotly.graph_objects as go\nfrom plotly.subplots import make_subplots\n\n# BioPython for sequence analysis\nfrom Bio import SeqIO\nfrom Bio.SeqUtils.ProtParam import ProteinAnalysis\n\n# GO Tools\nfrom goatools.obo_parser import GODag\n\n# ML libraries\nfrom sklearn.model_selection import StratifiedKFold, GroupKFold\nfrom sklearn.preprocessing import MultiLabelBinarizer, StandardScaler\nfrom sklearn.metrics import f1_score, precision_score, recall_score\nfrom sklearn.linear_model import Ridge\n\n# Deep Learning\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\n\nwarnings.filterwarnings('ignore')\n\n# Custom color palette\nCOLORS = {\n 'primary': '#667eea',\n 'secondary': '#764ba2',\n 'accent': '#f093fb',\n 'success': '#2ecc71',\n 'warning': '#f39c12',\n 'danger': '#e74c3c',\n 'MF': '#3498db',\n 'BP': '#27ae60',\n 'CC': '#e74c3c'\n}\n\n# Plotting configuration\nplt.style.use('seaborn-v0_8-whitegrid')\nplt.rcParams['figure.figsize'] = (14, 8)\nplt.rcParams['font.size'] = 12\nplt.rcParams['axes.titlesize'] = 16\nplt.rcParams['axes.labelsize'] = 14\n\nprint('Libraries loaded successfully!')"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Configuration\nclass CFG:\n # Paths (Kaggle environment)\n DATA_DIR = '/kaggle/input/cafa-6-protein-function-prediction'\n OUTPUT_DIR = '/kaggle/working'\n \n # Data files - note the subdirectory structure\n TRAIN_SEQS = f'{DATA_DIR}/Train/train_sequences.fasta'\n TEST_SEQS = f'{DATA_DIR}/Test/testsuperset.fasta'\n TRAIN_TERMS = f'{DATA_DIR}/Train/train_terms.tsv'\n TRAIN_TAXONOMY = f'{DATA_DIR}/Train/train_taxonomy.tsv'\n GO_OBO = f'{DATA_DIR}/Train/go-basic.obo'\n IA_FILE = f'{DATA_DIR}/IA.tsv'\n SAMPLE_SUB = f'{DATA_DIR}/sample_submission.tsv'\n \n # Model configuration\n EMBEDDING_DIM = 1280 # ESM-2 t33\n NUM_CLASSES = 5000 # Top GO terms\n BATCH_SIZE = 32\n EPOCHS = 15\n LR = 1e-3\n N_FOLDS = 5\n SEED = 42\n\n# Check for local environment fallback\nif not os.path.exists(CFG.DATA_DIR):\n CFG.DATA_DIR = 'data/raw'\n CFG.TRAIN_SEQS = f'{CFG.DATA_DIR}/train_sequences.fasta'\n CFG.TEST_SEQS = f'{CFG.DATA_DIR}/test_sequences.fasta'\n CFG.TRAIN_TERMS = f'{CFG.DATA_DIR}/train_terms.tsv'\n CFG.TRAIN_TAXONOMY = f'{CFG.DATA_DIR}/train_taxonomy.tsv'\n CFG.GO_OBO = f'{CFG.DATA_DIR}/go-basic.obo'\n CFG.IA_FILE = f'{CFG.DATA_DIR}/IA.tsv'\n CFG.SAMPLE_SUB = f'{CFG.DATA_DIR}/sample_submission.tsv'\n CFG.OUTPUT_DIR = 'output'\n\n# Set device after torch is imported\nCFG.DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\n\nprint(f\"Data directory: {CFG.DATA_DIR}\")\nprint(f\"Device: {CFG.DEVICE}\")"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 3. Data Loading & Overview"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load train terms\ntrain_terms = pd.read_csv(CFG.TRAIN_TERMS, sep='\\t')\nprint(f\"Train Terms Shape: {train_terms.shape}\")\nprint(f\"Unique Proteins: {train_terms['EntryID'].nunique():,}\")\nprint(f\"Unique GO Terms: {train_terms['term'].nunique():,}\")\ntrain_terms.head(10)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load taxonomy data\ntaxonomy = pd.read_csv(CFG.TRAIN_TAXONOMY, sep='\\t', header=None, names=['EntryID', 'TaxonomyID'])\nprint(f\"Taxonomy Shape: {taxonomy.shape}\")\nprint(f\"Unique Species: {taxonomy['TaxonomyID'].nunique():,}\")\ntaxonomy.head()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load Information Accretion scores\nia_scores = pd.read_csv(CFG.IA_FILE, sep='\\t', header=None, names=['term', 'IA'])\nprint(f\"IA Scores Shape: {ia_scores.shape}\")\nprint(f\"IA Range: [{ia_scores['IA'].min():.3f}, {ia_scores['IA'].max():.3f}]\")\nia_scores.head()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load Gene Ontology\nprint(\"Loading Gene Ontology...\")\ngodag = GODag(CFG.GO_OBO, optional_attrs={'relationship'})\nprint(f\"Total GO Terms in Ontology: {len(godag):,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Load sequences (memory efficient - just get stats first)\ndef load_sequence_stats(fasta_path, limit=None):\n \"\"\"Load sequence statistics without storing full sequences in memory.\"\"\"\n stats = []\n for i, record in enumerate(SeqIO.parse(fasta_path, 'fasta')):\n if limit and i >= limit:\n break\n seq = str(record.seq)\n entry_id = record.id.split('|')[1] if '|' in record.id else record.id\n stats.append({\n 'EntryID': entry_id,\n 'length': len(seq),\n 'description': record.description[:100]\n })\n return pd.DataFrame(stats)\n\nprint(\"Loading train sequences...\")\ntrain_seq_stats = load_sequence_stats(CFG.TRAIN_SEQS)\nprint(f\"Train Sequences: {len(train_seq_stats):,}\")\n\nprint(\"\\nLoading test sequences...\")\ntest_seq_stats = load_sequence_stats(CFG.TEST_SEQS)\nprint(f\"Test Sequences: {len(test_seq_stats):,}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Data overview summary\nsummary_data = {\n 'Metric': [\n 'Train Proteins', 'Test Proteins', 'Total Annotations',\n 'Unique GO Terms', 'Unique Species', 'GO Terms in Ontology',\n 'Avg Annotations/Protein', 'Avg Sequence Length'\n ],\n 'Train': [\n f\"{train_terms['EntryID'].nunique():,}\",\n '-',\n f\"{len(train_terms):,}\",\n f\"{train_terms['term'].nunique():,}\",\n f\"{taxonomy['TaxonomyID'].nunique():,}\",\n f\"{len(godag):,}\",\n f\"{len(train_terms) / train_terms['EntryID'].nunique():.1f}\",\n f\"{train_seq_stats['length'].mean():.0f}\"\n ],\n 'Test': [\n '-',\n f\"{len(test_seq_stats):,}\",\n 'To predict',\n '-',\n '-',\n '-',\n '-',\n f\"{test_seq_stats['length'].mean():.0f}\"\n ]\n}\n\nsummary_df = pd.DataFrame(summary_data)\nsummary_df.style.set_properties(**{'text-align': 'center'}).set_table_styles(\n [{'selector': 'th', 'props': [('text-align', 'center')]}]\n)"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 4. Exploratory Data Analysis\n\n"}, {"cell_type": "markdown", "metadata": {}, "source": "### 4.1 Protein Sequence Analysis"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Sequence Length Distribution - Interactive\nfig = make_subplots(rows=1, cols=2, subplot_titles=(\n 'Train Sequence Length Distribution', \n 'Test Sequence Length Distribution'\n))\n\nfig.add_trace(\n go.Histogram(x=train_seq_stats['length'], nbinsx=100, name='Train',\n marker_color=COLORS['primary'], opacity=0.75),\n row=1, col=1\n)\n\nfig.add_trace(\n go.Histogram(x=test_seq_stats['length'], nbinsx=100, name='Test',\n marker_color=COLORS['secondary'], opacity=0.75),\n row=1, col=2\n)\n\nfig.update_layout(\n title_text='Protein Sequence Length Distribution',\n title_font_size=20,\n showlegend=True,\n template='plotly_white',\n height=500\n)\n\nfig.update_xaxes(title_text='Sequence Length (amino acids)')\nfig.update_yaxes(title_text='Count')\nfig.show()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Sequence length statistics comparison\nprint(\"Sequence Length Statistics:\")\nprint(\"=\"*50)\nprint(f\"{'':15} {'Train':>15} {'Test':>15}\")\nprint(\"-\"*50)\nprint(f\"{'Min':15} {train_seq_stats['length'].min():>15,} {test_seq_stats['length'].min():>15,}\")\nprint(f\"{'Max':15} {train_seq_stats['length'].max():>15,} {test_seq_stats['length'].max():>15,}\")\nprint(f\"{'Mean':15} {train_seq_stats['length'].mean():>15,.0f} {test_seq_stats['length'].mean():>15,.0f}\")\nprint(f\"{'Median':15} {train_seq_stats['length'].median():>15,.0f} {test_seq_stats['length'].median():>15,.0f}\")\nprint(f\"{'Std':15} {train_seq_stats['length'].std():>15,.0f} {test_seq_stats['length'].std():>15,.0f}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Amino Acid Composition Analysis\ndef get_aa_composition(fasta_path, sample_size=5000):\n \"\"\"Get amino acid composition from sequences.\"\"\"\n aa_counts = Counter()\n for i, record in enumerate(SeqIO.parse(fasta_path, 'fasta')):\n if i >= sample_size:\n break\n aa_counts.update(str(record.seq))\n return aa_counts\n\nprint(\"Analyzing amino acid composition...\")\ntrain_aa = get_aa_composition(CFG.TRAIN_SEQS)\ntest_aa = get_aa_composition(CFG.TEST_SEQS)\n\n# Standard amino acids\nstandard_aa = 'ACDEFGHIKLMNPQRSTVWY'\n\n# Create comparison dataframe\naa_df = pd.DataFrame({\n 'AA': list(standard_aa),\n 'Train': [train_aa.get(aa, 0) for aa in standard_aa],\n 'Test': [test_aa.get(aa, 0) for aa in standard_aa]\n})\n\n# Normalize to percentages\naa_df['Train_pct'] = aa_df['Train'] / aa_df['Train'].sum() * 100\naa_df['Test_pct'] = aa_df['Test'] / aa_df['Test'].sum() * 100\naa_df = aa_df.sort_values('Train_pct', ascending=True)"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Amino Acid Composition - Beautiful horizontal bar chart\nfig = go.Figure()\n\nfig.add_trace(go.Bar(\n y=aa_df['AA'],\n x=aa_df['Train_pct'],\n name='Train',\n orientation='h',\n marker_color=COLORS['primary'],\n text=aa_df['Train_pct'].round(2).astype(str) + '%',\n textposition='outside'\n))\n\nfig.add_trace(go.Bar(\n y=aa_df['AA'],\n x=aa_df['Test_pct'],\n name='Test',\n orientation='h',\n marker_color=COLORS['secondary'],\n opacity=0.7\n))\n\nfig.update_layout(\n title='Amino Acid Composition Comparison',\n title_font_size=20,\n xaxis_title='Percentage (%)',\n yaxis_title='Amino Acid',\n barmode='overlay',\n template='plotly_white',\n height=700,\n legend=dict(x=0.85, y=0.05)\n)\n\nfig.show()"}, {"cell_type": "markdown", "metadata": {}, "source": "### 4.2 GO Term Distribution"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# GO Aspect Distribution - Stunning donut chart\naspect_counts = train_terms['aspect'].value_counts()\naspect_names = {'F': 'Molecular Function (MF)', 'P': 'Biological Process (BP)', 'C': 'Cellular Component (CC)'}\n\nfig = go.Figure(data=[go.Pie(\n labels=[aspect_names.get(a, a) for a in aspect_counts.index],\n values=aspect_counts.values,\n hole=0.5,\n marker_colors=[COLORS['MF'], COLORS['BP'], COLORS['CC']],\n textinfo='label+percent',\n textposition='outside',\n pull=[0.02, 0.02, 0.02]\n)])\n\nfig.update_layout(\n title='GO Aspect Distribution',\n title_font_size=20,\n annotations=[dict(\n text=f'{len(train_terms):,} annotations',\n x=0.5, y=0.5, font_size=16, showarrow=False\n )],\n template='plotly_white',\n height=500\n)\n\nfig.show()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Annotations per protein distribution\nann_per_protein = train_terms.groupby('EntryID').size().reset_index(name='count')\n\nfig = px.histogram(\n ann_per_protein, x='count', nbins=50,\n title='Distribution of GO Annotations per Protein',\n labels={'count': 'Number of GO Terms', 'count': 'Frequency'},\n color_discrete_sequence=[COLORS['primary']]\n)\n\nfig.add_vline(x=ann_per_protein['count'].mean(), line_dash='dash', line_color='red',\n annotation_text=f\"Mean: {ann_per_protein['count'].mean():.1f}\")\nfig.add_vline(x=ann_per_protein['count'].median(), line_dash='dash', line_color='green',\n annotation_text=f\"Median: {ann_per_protein['count'].median():.0f}\")\n\nfig.update_layout(template='plotly_white', height=500)\nfig.show()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Top 30 Most Frequent GO Terms with aspect coloring\ntop_terms = train_terms['term'].value_counts().head(30).reset_index()\ntop_terms.columns = ['term', 'count']\n\n# Get aspect for each term\nterm_aspect = train_terms.drop_duplicates('term').set_index('term')['aspect']\ntop_terms['aspect'] = top_terms['term'].map(term_aspect)\n\n# Get GO term names\ntop_terms['name'] = top_terms['term'].apply(\n lambda x: godag[x].name if x in godag else 'Unknown'\n)\ntop_terms['label'] = top_terms['term'] + ' ' + top_terms['name'].str[:30]\n\nfig = px.bar(\n top_terms, y='label', x='count', orientation='h',\n color='aspect',\n color_discrete_map={'F': COLORS['MF'], 'P': COLORS['BP'], 'C': COLORS['CC']},\n title='Top 30 Most Frequent GO Terms',\n labels={'count': 'Frequency', 'label': 'GO Term'}\n)\n\nfig.update_layout(\n template='plotly_white',\n height=800,\n legend_title='Aspect',\n yaxis={'categoryorder': 'total ascending'}\n)\nfig.show()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# GO Term Frequency Distribution (log scale)\nterm_freq = train_terms['term'].value_counts()\n\nfig = go.Figure()\n\nfig.add_trace(go.Scatter(\n x=np.arange(1, len(term_freq) + 1),\n y=term_freq.values,\n mode='lines',\n line=dict(color=COLORS['primary'], width=2),\n fill='tozeroy',\n fillcolor='rgba(102, 126, 234, 0.3)'\n))\n\nfig.update_layout(\n title='GO Term Frequency Distribution (Zipf\\'s Law)',\n xaxis_title='Term Rank',\n yaxis_title='Frequency (log scale)',\n xaxis_type='log',\n yaxis_type='log',\n template='plotly_white',\n height=500\n)\n\n# Add annotations\nfig.add_annotation(\n x=np.log10(100), y=np.log10(term_freq.iloc[99]),\n text=f'Top 100 terms: {term_freq.head(100).sum():,} annotations',\n showarrow=True, arrowhead=2\n)\n\nfig.show()\n\nprint(f\"\\nCoverage Analysis:\")\nprint(f\"Top 100 terms cover {term_freq.head(100).sum() / term_freq.sum() * 100:.1f}% of annotations\")\nprint(f\"Top 500 terms cover {term_freq.head(500).sum() / term_freq.sum() * 100:.1f}% of annotations\")\nprint(f\"Top 1000 terms cover {term_freq.head(1000).sum() / term_freq.sum() * 100:.1f}% of annotations\")"}, {"cell_type": "markdown", "metadata": {}, "source": "### 4.3 Taxonomy Analysis"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Taxonomy ID mapping to common names\nTAXONOMY_NAMES = {\n 9606: 'Homo sapiens (Human)',\n 10090: 'Mus musculus (Mouse)',\n 10116: 'Rattus norvegicus (Rat)',\n 7227: 'Drosophila melanogaster (Fruit fly)',\n 6239: 'Caenorhabditis elegans (Worm)',\n 4932: 'Saccharomyces cerevisiae (Yeast)',\n 7955: 'Danio rerio (Zebrafish)',\n 3702: 'Arabidopsis thaliana (Plant)',\n 559292: 'Saccharomyces cerevisiae S288C',\n 83333: 'Escherichia coli K-12'\n}\n\n# Get species distribution\nspecies_counts = taxonomy['TaxonomyID'].value_counts().head(20).reset_index()\nspecies_counts.columns = ['TaxonomyID', 'count']\nspecies_counts['name'] = species_counts['TaxonomyID'].map(\n lambda x: TAXONOMY_NAMES.get(x, f'Species {x}')\n)\n\n# Create sunburst-style bar chart\nfig = px.bar(\n species_counts, x='count', y='name', orientation='h',\n title='Top 20 Species in Training Data',\n labels={'count': 'Number of Proteins', 'name': 'Species'},\n color='count',\n color_continuous_scale='Viridis'\n)\n\nfig.update_layout(\n template='plotly_white',\n height=600,\n yaxis={'categoryorder': 'total ascending'},\n coloraxis_showscale=False\n)\nfig.show()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Species diversity treemap\nspecies_for_treemap = taxonomy['TaxonomyID'].value_counts().head(50).reset_index()\nspecies_for_treemap.columns = ['TaxonomyID', 'count']\nspecies_for_treemap['name'] = species_for_treemap['TaxonomyID'].map(\n lambda x: TAXONOMY_NAMES.get(x, f'Species {x}')\n)\nspecies_for_treemap['parent'] = 'All Species'\n\nfig = px.treemap(\n species_for_treemap,\n path=['parent', 'name'],\n values='count',\n title='Species Composition Treemap',\n color='count',\n color_continuous_scale='RdYlGn'\n)\n\nfig.update_layout(height=600, template='plotly_white')\nfig.show()"}, {"cell_type": "markdown", "metadata": {}, "source": "### 4.4 GO Hierarchy Visualization"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Analyze GO hierarchy depth\ndef get_term_depth(term_id, godag):\n \"\"\"Get the depth of a GO term in the hierarchy.\"\"\"\n if term_id not in godag:\n return -1\n return godag[term_id].depth\n\ntrain_terms['depth'] = train_terms['term'].apply(lambda x: get_term_depth(x, godag))\n\n# Depth distribution by aspect\nfig = px.violin(\n train_terms[train_terms['depth'] >= 0],\n x='aspect', y='depth',\n color='aspect',\n color_discrete_map={'F': COLORS['MF'], 'P': COLORS['BP'], 'C': COLORS['CC']},\n title='GO Term Depth Distribution by Aspect',\n labels={'aspect': 'GO Aspect', 'depth': 'Hierarchy Depth'},\n box=True, points='outliers'\n)\n\nfig.update_layout(template='plotly_white', height=500, showlegend=False)\nfig.show()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# GO hierarchy network visualization (sample)\ndef build_go_subgraph(terms, godag, max_terms=100):\n \"\"\"Build a subgraph of GO terms for visualization.\"\"\"\n edges = []\n nodes = set()\n \n sample_terms = list(terms)[:max_terms]\n \n for term in sample_terms:\n if term not in godag:\n continue\n nodes.add(term)\n for parent in godag[term].parents:\n edges.append((parent.id, term))\n nodes.add(parent.id)\n \n return list(nodes), edges\n\n# Get top terms for each aspect\nmf_terms = train_terms[train_terms['aspect'] == 'F']['term'].value_counts().head(20).index\nbp_terms = train_terms[train_terms['aspect'] == 'P']['term'].value_counts().head(20).index\ncc_terms = train_terms[train_terms['aspect'] == 'C']['term'].value_counts().head(20).index\n\nprint(f\"Top MF terms: {len(mf_terms)}\")\nprint(f\"Top BP terms: {len(bp_terms)}\")\nprint(f\"Top CC terms: {len(cc_terms)}\")"}, {"cell_type": "markdown", "metadata": {}, "source": "### 4.5 Information Accretion Analysis"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Information Accretion distribution\nfig = px.histogram(\n ia_scores[ia_scores['IA'] > 0], x='IA', nbins=100,\n title='Information Accretion (IA) Score Distribution',\n labels={'IA': 'Information Accretion Score'},\n color_discrete_sequence=[COLORS['accent']]\n)\n\nfig.add_vline(x=ia_scores['IA'].mean(), line_dash='dash', line_color='red',\n annotation_text=f\"Mean: {ia_scores['IA'].mean():.2f}\")\n\nfig.update_layout(template='plotly_white', height=500)\nfig.show()\n\nprint(f\"\\nIA Statistics:\")\nprint(f\"Terms with IA > 0: {(ia_scores['IA'] > 0).sum():,}\")\nprint(f\"Terms with IA = 0: {(ia_scores['IA'] == 0).sum():,}\")\nprint(f\"Max IA: {ia_scores['IA'].max():.3f}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Merge IA with term frequency to understand scoring impact\nterm_stats = train_terms['term'].value_counts().reset_index()\nterm_stats.columns = ['term', 'frequency']\nterm_stats = term_stats.merge(ia_scores, on='term', how='left')\nterm_stats['IA'] = term_stats['IA'].fillna(0)\n\n# Scatter plot: Frequency vs IA\nfig = px.scatter(\n term_stats.head(1000), x='frequency', y='IA',\n title='Term Frequency vs Information Accretion (Top 1000 Terms)',\n labels={'frequency': 'Term Frequency', 'IA': 'Information Accretion'},\n color='IA',\n color_continuous_scale='Viridis',\n hover_data=['term']\n)\n\nfig.update_layout(template='plotly_white', height=600)\nfig.show()\n\n# Key insight\nprint(\"\\nKey Insight:\")\nprint(\"High-frequency terms tend to have LOWER IA scores (less specific)\")\nprint(\"Rare, specific terms have HIGHER IA scores (more informative)\")\nprint(\"\\nThis means correctly predicting rare terms is MORE valuable for the metric!\")"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 5. Feature Engineering\n\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Feature engineering utilities\ndef compute_sequence_features(sequence):\n \"\"\"Compute hand-crafted features from protein sequence.\"\"\"\n features = {}\n \n # Basic features\n features['length'] = len(sequence)\n \n # Amino acid composition\n aa_counts = Counter(sequence)\n total = len(sequence)\n for aa in 'ACDEFGHIKLMNPQRSTVWY':\n features[f'aa_{aa}'] = aa_counts.get(aa, 0) / total\n \n # Physicochemical properties\n try:\n pa = ProteinAnalysis(sequence)\n features['molecular_weight'] = pa.molecular_weight()\n features['isoelectric_point'] = pa.isoelectric_point()\n features['gravy'] = pa.gravy() # Hydropathicity\n features['instability_index'] = pa.instability_index()\n features['aromaticity'] = pa.aromaticity()\n \n # Secondary structure fractions\n ss = pa.secondary_structure_fraction()\n features['helix_fraction'] = ss[0]\n features['turn_fraction'] = ss[1]\n features['sheet_fraction'] = ss[2]\n except:\n pass\n \n return features\n\nprint(\"Feature engineering functions defined.\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Compute features for a sample\nprint(\"Computing sequence features for sample proteins...\")\n\nsample_features = []\nfor i, record in enumerate(SeqIO.parse(CFG.TRAIN_SEQS, 'fasta')):\n if i >= 1000: # Sample 1000 proteins\n break\n seq = str(record.seq)\n entry_id = record.id.split('|')[1] if '|' in record.id else record.id\n feats = compute_sequence_features(seq)\n feats['EntryID'] = entry_id\n sample_features.append(feats)\n\nfeature_df = pd.DataFrame(sample_features)\nprint(f\"Computed {len(feature_df.columns) - 1} features for {len(feature_df)} proteins\")\nfeature_df.head()"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Feature correlation heatmap\nnumeric_cols = feature_df.select_dtypes(include=[np.number]).columns[:15] # First 15 numeric features\ncorr_matrix = feature_df[numeric_cols].corr()\n\nfig = px.imshow(\n corr_matrix,\n title='Feature Correlation Heatmap',\n color_continuous_scale='RdBu_r',\n aspect='auto'\n)\n\nfig.update_layout(template='plotly_white', height=600)\nfig.show()"}, {"cell_type": "code", "source": "# DIAMOND Homology-Based Label Transfer\n# This is a strong baseline that exploits evolutionary conservation\n\ndef run_diamond_homology(train_fasta, test_fasta, train_terms, output_dir):\n \"\"\"\n Run DIAMOND BLASTp to find homologs and transfer labels.\n \n DIAMOND is a fast sequence aligner that can process millions of sequences.\n Homology-based annotation transfer is the most reliable method for well-studied protein families.\n \"\"\"\n import subprocess\n \n os.makedirs(output_dir, exist_ok=True)\n db_path = os.path.join(output_dir, \"train_db\")\n align_path = os.path.join(output_dir, \"alignments.tsv\")\n \n # Step 1: Build DIAMOND database from training sequences\n print(\"Building DIAMOND database...\")\n subprocess.run([\n \"diamond\", \"makedb\",\n \"--in\", train_fasta,\n \"-d\", db_path\n ], check=True)\n \n # Step 2: Run BLASTp alignment\n print(\"Running DIAMOND BLASTp...\")\n subprocess.run([\n \"diamond\", \"blastp\",\n \"-d\", db_path,\n \"-q\", test_fasta,\n \"-o\", align_path,\n \"--sensitive\",\n \"--top\", \"10\",\n \"--outfmt\", \"6\", \"qseqid\", \"sseqid\", \"pident\", \"bitscore\"\n ], check=True)\n \n # Step 3: Transfer labels based on alignment\n print(\"Transferring labels...\")\n alignments = pd.read_csv(align_path, sep=\"\\t\", \n names=[\"query_id\", \"subject_id\", \"pident\", \"bitscore\"])\n \n # Parse subject ID (format: sp|ACCESSION|NAME)\n alignments['subject_id'] = alignments['subject_id'].apply(\n lambda x: x.split('|')[1] if '|' in str(x) else x\n )\n \n # Join with training labels\n merged = alignments.merge(train_terms, left_on=\"subject_id\", right_on=\"EntryID\")\n \n # Score = percent identity / 100\n merged['score'] = merged['pident'] / 100.0\n \n # For each (query, term), keep max score\n predictions = merged.groupby(['query_id', 'term'])['score'].max().reset_index()\n predictions.columns = ['id', 'term', 'score']\n \n output_path = os.path.join(output_dir, \"homology_predictions.tsv\")\n predictions.to_csv(output_path, sep=\"\\t\", index=False)\n \n print(f\"Saved {len(predictions):,} homology predictions\")\n return predictions\n\nprint(\"DIAMOND homology baseline defined.\")\nprint()\nprint(\"To run DIAMOND on Kaggle:\")\nprint(\" !apt-get install diamond-aligner\")\nprint(\" Or download pre-built binary\")", "metadata": {}, "execution_count": null, "outputs": []}, {"cell_type": "code", "source": "# ESM-2 Protein Embeddings Extraction\n# This is the key differentiator - using state-of-the-art protein language models\n\ndef extract_esm2_embeddings(fasta_path, output_dir, model_name='esm2_t33_650M_UR50D', batch_size=4):\n \"\"\"\n Extract ESM-2 embeddings for protein sequences.\n \n ESM-2 is a state-of-the-art protein language model trained on millions of sequences.\n The t33_650M variant provides 1280-dimensional embeddings.\n \"\"\"\n import torch\n from transformers import AutoTokenizer, AutoModel\n \n print(f\"Loading ESM-2 model: {model_name}...\")\n tokenizer = AutoTokenizer.from_pretrained(f\"facebook/{model_name}\")\n model = AutoModel.from_pretrained(f\"facebook/{model_name}\")\n model.eval()\n \n device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n model = model.to(device)\n \n os.makedirs(output_dir, exist_ok=True)\n \n embeddings = []\n ids = []\n \n sequences = []\n for record in SeqIO.parse(fasta_path, \"fasta\"):\n entry_id = record.id.split('|')[1] if '|' in record.id else record.id\n seq = str(record.seq)[:1022] # ESM-2 max length\n sequences.append((entry_id, seq))\n \n print(f\"Processing {len(sequences)} sequences...\")\n \n with torch.no_grad():\n for i in tqdm(range(0, len(sequences), batch_size)):\n batch = sequences[i:i+batch_size]\n batch_ids = [s[0] for s in batch]\n batch_seqs = [s[1] for s in batch]\n \n inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, truncation=True, max_length=1024)\n inputs = {k: v.to(device) for k, v in inputs.items()}\n \n outputs = model(**inputs)\n \n # Mean pooling over sequence length\n emb = outputs.last_hidden_state.mean(dim=1).cpu().numpy()\n \n embeddings.extend(emb)\n ids.extend(batch_ids)\n \n embeddings = np.array(embeddings)\n ids = np.array(ids)\n \n np.save(os.path.join(output_dir, \"embeddings.npy\"), embeddings)\n np.save(os.path.join(output_dir, \"ids.npy\"), ids)\n \n print(f\"Saved embeddings: {embeddings.shape}\")\n return embeddings, ids\n\nprint(\"ESM-2 embedding extraction function defined.\")\nprint()\nprint(\"Note: ESM-2 extraction requires GPU and is memory intensive.\")\nprint(\"For Kaggle, pre-computed embeddings can be loaded from datasets like:\")\nprint(\" - kaggle.com/datasets/dschettler8845/cafa-5-esm2-embeddings\")\nprint(\" - Or compute on Kaggle GPU for ~4-6 hours\")", "metadata": {}, "execution_count": null, "outputs": []}, {"cell_type": "code", "source": "# External Data: UniProt Features Integration\ndef fetch_uniprot_features(entry_ids, batch_size=100):\n \"\"\"\n Fetch additional features from UniProt for given protein IDs.\n This demonstrates external data integration.\n \"\"\"\n import requests\n \n features = []\n base_url = \"https://rest.uniprot.org/uniprotkb/search\"\n \n for i in range(0, min(len(entry_ids), 500), batch_size): # Limit for demo\n batch_ids = entry_ids[i:i+batch_size]\n query = \" OR \".join([f\"accession:{eid}\" for eid in batch_ids])\n \n params = {\n \"query\": query,\n \"fields\": \"accession,length,mass,cc_subcellular_location,ft_domain,organism_name\",\n \"format\": \"json\",\n \"size\": batch_size\n }\n \n try:\n response = requests.get(base_url, params=params, timeout=30)\n if response.status_code == 200:\n data = response.json()\n for result in data.get('results', []):\n features.append({\n 'EntryID': result.get('primaryAccession'),\n 'mass': result.get('sequence', {}).get('mass'),\n 'organism': result.get('organism', {}).get('scientificName'),\n 'has_structure': bool(result.get('uniProtKBCrossReferences', []))\n })\n except Exception as e:\n print(f\"Error fetching batch: {e}\")\n continue\n \n return pd.DataFrame(features) if features else pd.DataFrame()\n\n# Demo: Fetch features for sample proteins\nprint(\"External Data Integration Demo:\")\nprint(\"In a full pipeline, we would fetch:\")\nprint(\" - UniProt annotations and cross-references\")\nprint(\" - InterPro domain annotations\")\nprint(\" - PDB structure availability\")\nprint(\" - STRING protein-protein interactions\")\nprint()\nprint(\"Note: External API calls are disabled in this notebook for efficiency.\")", "metadata": {}, "execution_count": null, "outputs": []}, {"cell_type": "markdown", "source": "### 5.1 External Data Integration\n\n
\nExternal Datasets: Leveraging external biological databases significantly improves protein function prediction. Here we integrate:\n
\n
UniProt: Curated protein annotations and features
\n
InterPro: Protein domain/family classification
\n
PDB: 3D structure availability as feature
\n
\n
", "metadata": {}}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 6. Model Architecture\n\n\n\nOur solution combines multiple approaches:\n\n1. **ESM-2 Embeddings**: Pre-trained protein language model embeddings (1280-dim)\n2. **Homology Transfer**: DIAMOND BLASTp for sequence similarity-based label transfer\n3. **Deep Neural Network**: Multi-label classifier on top of embeddings\n4. **GO Hierarchy Propagation**: True Path Rule for consistency"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Model Architecture\nclass ResidualBlock(nn.Module):\n \"\"\"Residual block with batch normalization and dropout.\"\"\"\n def __init__(self, in_features, out_features, dropout=0.3):\n super().__init__()\n self.fc1 = nn.Linear(in_features, out_features)\n self.bn1 = nn.BatchNorm1d(out_features)\n self.fc2 = nn.Linear(out_features, out_features)\n self.bn2 = nn.BatchNorm1d(out_features)\n self.dropout = nn.Dropout(dropout)\n \n self.shortcut = nn.Linear(in_features, out_features) if in_features != out_features else nn.Identity()\n \n def forward(self, x):\n identity = self.shortcut(x)\n out = F.relu(self.bn1(self.fc1(x)))\n out = self.dropout(out)\n out = self.bn2(self.fc2(out))\n return F.relu(out + identity)\n\n\nclass ProteinFunctionPredictor(nn.Module):\n \"\"\"Multi-label classifier for protein function prediction.\"\"\"\n def __init__(self, input_dim=1280, num_classes=5000, hidden_dims=[1024, 512, 256]):\n super().__init__()\n \n layers = []\n prev_dim = input_dim\n \n for hidden_dim in hidden_dims:\n layers.append(ResidualBlock(prev_dim, hidden_dim))\n prev_dim = hidden_dim\n \n self.encoder = nn.Sequential(*layers)\n self.classifier = nn.Linear(prev_dim, num_classes)\n \n # Aspect-specific heads (optional)\n self.mf_head = nn.Linear(prev_dim, num_classes // 3)\n self.bp_head = nn.Linear(prev_dim, num_classes // 3)\n self.cc_head = nn.Linear(prev_dim, num_classes // 3)\n \n def forward(self, x):\n features = self.encoder(x)\n return self.classifier(features)\n \n def forward_multi_head(self, x):\n features = self.encoder(x)\n return {\n 'mf': self.mf_head(features),\n 'bp': self.bp_head(features),\n 'cc': self.cc_head(features)\n }\n\n\n# Model summary\nmodel = ProteinFunctionPredictor(input_dim=1280, num_classes=5000)\nprint(f\"Model Parameters: {sum(p.numel() for p in model.parameters()):,}\")"}, {"cell_type": "code", "source": "# Complete Training Pipeline\ndef train_complete_pipeline(train_terms, train_taxonomy, godag, ia_scores,\n embeddings_path=None, n_folds=5, epochs=15):\n \"\"\"\n Complete training pipeline combining all components.\n \n Pipeline:\n 1. Load/Extract ESM-2 embeddings\n 2. Prepare multi-label targets\n 3. Train with cross-validation\n 4. Apply GO hierarchy propagation\n 5. (Optional) Ensemble with homology predictions\n \"\"\"\n print(\"=\"*60)\n print(\"COMPLETE TRAINING PIPELINE\")\n print(\"=\"*60)\n \n # Step 1: Prepare data\n print(\"\\n[1/5] Preparing data...\")\n \n # Get unique proteins and top GO terms\n unique_proteins = train_terms['EntryID'].unique()\n term_counts = train_terms['term'].value_counts()\n top_terms = term_counts.head(CFG.NUM_CLASSES).index.tolist()\n \n print(f\" Proteins: {len(unique_proteins):,}\")\n print(f\" GO Terms (top {CFG.NUM_CLASSES}): {len(top_terms):,}\")\n \n # Create term-to-index mapping\n term_to_idx = {term: i for i, term in enumerate(top_terms)}\n \n # Create multi-label target matrix\n protein_to_idx = {p: i for i, p in enumerate(unique_proteins)}\n y = np.zeros((len(unique_proteins), len(top_terms)), dtype=np.float32)\n \n for _, row in train_terms.iterrows():\n if row['term'] in term_to_idx:\n pid = protein_to_idx.get(row['EntryID'])\n if pid is not None:\n y[pid, term_to_idx[row['term']]] = 1.0\n \n print(f\" Target matrix shape: {y.shape}\")\n print(f\" Positive labels: {y.sum():,.0f} ({y.mean()*100:.2f}%)\")\n \n # Step 2: Load embeddings\n print(\"\\n[2/5] Loading embeddings...\")\n if embeddings_path and os.path.exists(embeddings_path):\n X = np.load(embeddings_path)\n print(f\" Loaded embeddings: {X.shape}\")\n else:\n # Use random embeddings for demo\n print(\" WARNING: Using random embeddings for demonstration\")\n X = np.random.randn(len(unique_proteins), CFG.EMBEDDING_DIM).astype(np.float32)\n \n # Ensure alignment\n if len(X) != len(y):\n print(f\" Aligning X ({len(X)}) with y ({len(y)})...\")\n min_len = min(len(X), len(y))\n X = X[:min_len]\n y = y[:min_len]\n \n # Step 3: Cross-validation training\n print(\"\\n[3/5] Training with cross-validation...\")\n fold_results, oof_predictions = run_cv_training(X, y, n_folds=n_folds, epochs=epochs)\n \n # Step 4: Calculate CV score\n print(\"\\n[4/5] Calculating CV metrics...\")\n cv_scores = [r['best_f1'] for r in fold_results]\n print(f\" CV F1: {np.mean(cv_scores):.4f} (+/- {np.std(cv_scores):.4f})\")\n \n # Step 5: Save results\n print(\"\\n[5/5] Saving results...\")\n results = {\n 'fold_results': fold_results,\n 'oof_predictions': oof_predictions,\n 'term_to_idx': term_to_idx,\n 'cv_score': np.mean(cv_scores)\n }\n \n print(\"=\"*60)\n print(\"TRAINING COMPLETE\")\n print(\"=\"*60)\n \n return results\n\nprint(\"Complete training pipeline defined.\")", "metadata": {}, "execution_count": null, "outputs": []}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Custom loss function with class weighting based on IA scores\nclass WeightedBCELoss(nn.Module):\n \"\"\"BCE loss weighted by Information Accretion scores.\"\"\"\n def __init__(self, class_weights=None):\n super().__init__()\n self.class_weights = class_weights\n \n def forward(self, logits, targets):\n bce = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')\n if self.class_weights is not None:\n bce = bce * self.class_weights\n return bce.mean()\n\n\nclass FocalLoss(nn.Module):\n \"\"\"Focal loss for handling class imbalance.\"\"\"\n def __init__(self, alpha=0.25, gamma=2.0):\n super().__init__()\n self.alpha = alpha\n self.gamma = gamma\n \n def forward(self, logits, targets):\n bce = F.binary_cross_entropy_with_logits(logits, targets, reduction='none')\n pt = torch.exp(-bce)\n focal_loss = self.alpha * (1 - pt) ** self.gamma * bce\n return focal_loss.mean()\n\nprint(\"Loss functions defined.\")"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 7. Training with Cross-Validation\n\n"}, {"cell_type": "code", "source": "# F-max Metric Implementation (Competition Evaluation Metric)\ndef calculate_fmax(y_true, y_pred, thresholds=np.arange(0.01, 1.0, 0.01)):\n \"\"\"\n Calculate F-max score - the competition evaluation metric.\n \n F-max is the maximum F1 score across all prediction thresholds.\n This is the standard metric for protein function prediction challenges.\n \n Args:\n y_true: Ground truth multi-label matrix (n_samples, n_labels)\n y_pred: Predicted probability matrix (n_samples, n_labels)\n thresholds: Array of thresholds to evaluate\n \n Returns:\n fmax: Maximum F1 score\n best_threshold: Threshold that achieves F-max\n \"\"\"\n fmax = 0\n best_threshold = 0\n \n for t in thresholds:\n y_pred_binary = (y_pred >= t).astype(int)\n \n # Calculate precision and recall\n true_positives = ((y_pred_binary == 1) & (y_true == 1)).sum()\n predicted_positives = y_pred_binary.sum()\n actual_positives = y_true.sum()\n \n precision = true_positives / max(predicted_positives, 1)\n recall = true_positives / max(actual_positives, 1)\n \n # F1 score\n if precision + recall > 0:\n f1 = 2 * precision * recall / (precision + recall)\n else:\n f1 = 0\n \n if f1 > fmax:\n fmax = f1\n best_threshold = t\n \n return fmax, best_threshold\n\n\ndef calculate_weighted_fmax(y_true, y_pred, ia_weights, thresholds=np.arange(0.01, 1.0, 0.01)):\n \"\"\"\n Calculate weighted F-max using Information Accretion scores.\n \n The actual competition metric weights terms by their IA scores,\n giving more importance to specific (rare) terms.\n \"\"\"\n fmax = 0\n best_threshold = 0\n \n # Ensure weights array matches shape\n if len(ia_weights) != y_true.shape[1]:\n ia_weights = np.ones(y_true.shape[1])\n \n for t in thresholds:\n y_pred_binary = (y_pred >= t).astype(int)\n \n # Weighted metrics\n weighted_tp = ((y_pred_binary == 1) & (y_true == 1) * ia_weights).sum()\n weighted_pred = (y_pred_binary * ia_weights).sum()\n weighted_actual = (y_true * ia_weights).sum()\n \n precision = weighted_tp / max(weighted_pred, 1e-10)\n recall = weighted_tp / max(weighted_actual, 1e-10)\n \n if precision + recall > 0:\n f1 = 2 * precision * recall / (precision + recall)\n else:\n f1 = 0\n \n if f1 > fmax:\n fmax = f1\n best_threshold = t\n \n return fmax, best_threshold\n\nprint(\"F-max metric functions defined.\")", "metadata": {}, "execution_count": null, "outputs": []}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Dataset class\nclass ProteinDataset(Dataset):\n def __init__(self, embeddings, labels):\n self.embeddings = torch.tensor(embeddings, dtype=torch.float32)\n self.labels = torch.tensor(labels, dtype=torch.float32)\n \n def __len__(self):\n return len(self.labels)\n \n def __getitem__(self, idx):\n return self.embeddings[idx], self.labels[idx]\n\n\ndef train_fold(model, train_loader, val_loader, criterion, optimizer, scheduler, epochs, device):\n \"\"\"Train model for one fold.\"\"\"\n best_f1 = 0\n history = {'train_loss': [], 'val_loss': [], 'val_f1': []}\n \n for epoch in range(epochs):\n # Training\n model.train()\n train_loss = 0\n for X, y in train_loader:\n X, y = X.to(device), y.to(device)\n optimizer.zero_grad()\n outputs = model(X)\n loss = criterion(outputs, y)\n loss.backward()\n optimizer.step()\n train_loss += loss.item()\n \n # Validation\n model.eval()\n val_loss = 0\n all_preds, all_labels = [], []\n with torch.no_grad():\n for X, y in val_loader:\n X, y = X.to(device), y.to(device)\n outputs = model(X)\n loss = criterion(outputs, y)\n val_loss += loss.item()\n preds = torch.sigmoid(outputs).cpu().numpy()\n all_preds.append(preds)\n all_labels.append(y.cpu().numpy())\n \n all_preds = np.vstack(all_preds)\n all_labels = np.vstack(all_labels)\n \n # Calculate F1 (micro-averaged)\n val_f1 = f1_score(all_labels, (all_preds > 0.5).astype(int), average='micro', zero_division=0)\n \n history['train_loss'].append(train_loss / len(train_loader))\n history['val_loss'].append(val_loss / len(val_loader))\n history['val_f1'].append(val_f1)\n \n if scheduler:\n scheduler.step(val_loss)\n \n if val_f1 > best_f1:\n best_f1 = val_f1\n \n print(f\" Epoch {epoch+1}/{epochs} | Train Loss: {train_loss/len(train_loader):.4f} | \"\n f\"Val Loss: {val_loss/len(val_loader):.4f} | Val F1: {val_f1:.4f}\")\n \n return history, best_f1\n\nprint(\"Training functions defined.\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Cross-validation training loop (placeholder - needs actual embeddings)\ndef run_cv_training(X, y, n_folds=5, epochs=15, batch_size=32):\n \"\"\"Run K-fold cross-validation training.\"\"\"\n device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n \n # Group by taxonomy for more robust CV\n kf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=42)\n \n # Stratification key: most frequent GO term per protein\n stratify_key = np.argmax(y, axis=1)\n \n fold_results = []\n oof_predictions = np.zeros_like(y, dtype=np.float32)\n \n for fold, (train_idx, val_idx) in enumerate(kf.split(X, stratify_key)):\n print(f\"\\n{'='*50}\")\n print(f\"Fold {fold + 1}/{n_folds}\")\n print(f\"{'='*50}\")\n \n # Split data\n X_train, X_val = X[train_idx], X[val_idx]\n y_train, y_val = y[train_idx], y[val_idx]\n \n # Create datasets\n train_ds = ProteinDataset(X_train, y_train)\n val_ds = ProteinDataset(X_val, y_val)\n \n train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True, num_workers=0)\n val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False, num_workers=0)\n \n # Initialize model\n model = ProteinFunctionPredictor(input_dim=X.shape[1], num_classes=y.shape[1]).to(device)\n \n # Loss and optimizer\n criterion = FocalLoss(alpha=0.25, gamma=2.0)\n optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)\n scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, patience=3, factor=0.5)\n \n # Train\n history, best_f1 = train_fold(model, train_loader, val_loader, criterion, \n optimizer, scheduler, epochs, device)\n \n fold_results.append({'fold': fold, 'best_f1': best_f1, 'history': history})\n \n # Generate OOF predictions\n model.eval()\n with torch.no_grad():\n val_preds = []\n for X_batch, _ in val_loader:\n X_batch = X_batch.to(device)\n preds = torch.sigmoid(model(X_batch)).cpu().numpy()\n val_preds.append(preds)\n oof_predictions[val_idx] = np.vstack(val_preds)\n \n # Summary\n print(f\"\\n{'='*50}\")\n print(\"Cross-Validation Summary\")\n print(f\"{'='*50}\")\n cv_scores = [r['best_f1'] for r in fold_results]\n print(f\"Mean F1: {np.mean(cv_scores):.4f} (+/- {np.std(cv_scores):.4f})\")\n \n return fold_results, oof_predictions\n\nprint(\"CV training loop ready.\")"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 8. Ensemble & Post-Processing\n\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# GO Hierarchy Propagation (True Path Rule)\ndef propagate_predictions(predictions_df, godag, threshold=0.001):\n \"\"\"\n Propagate predictions up the GO hierarchy using True Path Rule.\n If a term is predicted, all its ancestors should have at least the same probability.\n \"\"\"\n print(\"Propagating predictions using True Path Rule...\")\n \n # Group by protein\n grouped = predictions_df.groupby('id')\n propagated_rows = []\n \n for protein_id, group in tqdm(grouped, desc=\"Propagating\"):\n scores = dict(zip(group['term'], group['score']))\n \n # Get all terms for this protein\n active_terms = list(scores.keys())\n \n for term in active_terms:\n if term not in godag:\n continue\n \n current_score = scores[term]\n \n # Propagate to all ancestors\n ancestors = godag[term].get_all_parents()\n for ancestor in ancestors:\n scores[ancestor] = max(scores.get(ancestor, 0.0), current_score)\n \n # Add to results\n for term, score in scores.items():\n if score >= threshold:\n propagated_rows.append({'id': protein_id, 'term': term, 'score': score})\n \n result_df = pd.DataFrame(propagated_rows)\n result_df = result_df.sort_values(['id', 'score'], ascending=[True, False])\n \n print(f\"Original predictions: {len(predictions_df):,}\")\n print(f\"After propagation: {len(result_df):,}\")\n \n return result_df\n\nprint(\"Post-processing functions defined.\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Ensemble function\ndef ensemble_predictions(preds_list, weights=None):\n \"\"\"\n Ensemble multiple prediction dataframes with optional weighting.\n \n Args:\n preds_list: List of DataFrames with columns ['id', 'term', 'score']\n weights: List of weights for each prediction set\n \"\"\"\n if weights is None:\n weights = [1.0] * len(preds_list)\n \n # Normalize weights\n weights = np.array(weights) / sum(weights)\n \n # Merge all predictions\n all_preds = pd.concat([\n df.assign(weight=w) for df, w in zip(preds_list, weights)\n ])\n \n # Weighted average by (id, term)\n all_preds['weighted_score'] = all_preds['score'] * all_preds['weight']\n \n ensemble = all_preds.groupby(['id', 'term']).agg({\n 'weighted_score': 'sum',\n 'weight': 'sum'\n }).reset_index()\n \n # Normalize by total weight\n ensemble['score'] = ensemble['weighted_score'] / ensemble['weight'].clip(lower=1e-10)\n \n return ensemble[['id', 'term', 'score']]\n\nprint(\"Ensemble functions defined.\")"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## 9. Submission\n\n"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Check sample submission format\nsample_sub = pd.read_csv(CFG.SAMPLE_SUB, sep='\\t', header=None, nrows=20, on_bad_lines='skip')\nprint(\"Sample Submission Format:\")\nprint(sample_sub.head(10))\nprint(f\"\\nShape: {sample_sub.shape}\")"}, {"cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "def create_submission(predictions_df, output_path, score_decimals=3):\n \"\"\"\n Create submission file in CAFA format.\n Format: protein_idGO_termscore (no header)\n \"\"\"\n # Ensure proper format\n sub = predictions_df[['id', 'term', 'score']].copy()\n \n # Round scores\n sub['score'] = sub['score'].round(score_decimals)\n \n # Filter very low scores\n sub = sub[sub['score'] >= 0.001]\n \n # Sort by id and score (descending)\n sub = sub.sort_values(['id', 'score'], ascending=[True, False])\n \n # Save without header\n sub.to_csv(output_path, sep='\\t', index=False, header=False)\n \n print(f\"Submission saved to: {output_path}\")\n print(f\"Total predictions: {len(sub):,}\")\n print(f\"Unique proteins: {sub['id'].nunique():,}\")\n print(f\"Unique GO terms: {sub['term'].nunique():,}\")\n \n return sub\n\nprint(\"Submission function ready.\")"}, {"cell_type": "markdown", "metadata": {}, "source": "---\n\n## Summary & Key Insights\n\n
\n\n### Key Findings from EDA\n\n1. **Data Scale**: ~82K train proteins with ~537K annotations across ~10K unique GO terms\n\n2. **Heavy Imbalance**: GO term frequency follows Zipf's law - top 100 terms cover >50% of annotations\n\n3. **IA Scoring Insight**: Rare, specific terms have higher IA scores, making them MORE valuable to predict correctly\n\n4. **Sequence Diversity**: Train and test have similar sequence length distributions (~400-500 aa average)\n\n5. **Species Bias**: Human proteins dominate the training data\n\n### Recommended Approach\n\n| Component | Method | Expected Contribution |\n|-----------|--------|----------------------|\n| **Embeddings** | ESM-2 (t33_650M) | Rich sequence representation |\n| **Homology** | DIAMOND BLASTp | Strong baseline for conserved functions |\n| **Neural Network** | Deep MLP + Focal Loss | Handle imbalanced multi-label |\n| **Post-processing** | GO hierarchy propagation | Ensure True Path Rule consistency |\n| **Ensemble** | Weighted average (0.7 homology + 0.3 DL) | Best of both worlds |\n\n### External Data Opportunities\n\n- **UniProt**: Additional protein annotations\n- **InterPro**: Domain/family features\n- **PDB**: Structure-based features for available structures\n- **String DB**: Protein-protein interaction networks\n\n