{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","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":118765,"databundleVersionId":15231210,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Stanford RNA 3D Folding Part 2: TBM Approach\n\n## Inspired by Part 1 Competition Winners\n\n---\n\n### Part 1 Competition Insights\n\nThe **Stanford RNA 3D Folding** competition (Part 1) revealed that **Template-Based Modeling (TBM)** was highly effective for RNA structure prediction.\n\n**Key findings from Part 1 discussions:**\n- TBM approaches dominated the leaderboard\n- The 1st place solution used a hybrid TBM + DRfold2 approach\n- Pure TBM pipelines achieved competitive scores\n- No extensive GPU training was required for top performance\n\n### Key Strategies from Top Solutions\n\n1. **5-Step TBM Pipeline**:\n   - **Search**: Find candidate templates using similarity scoring\n   - **Alignment**: Use established aligners (e.g., BioPython)\n   - **Transfer**: Map template coordinates to query\n   - **Gap Fill**: Fill missing positions (~5.9Å between C1' atoms)\n   - **Refinement**: Local optimization to reduce clashes\n\n2. **Composite Similarity Scoring**:\n   ```\n   Score = 0.4 × Global + 0.3 × Local + 0.2 × Feature + 0.1 × Kmer\n   ```\n\n3. **Template Library Size Matters**: More templates = better coverage\n\n4. **No GPU Training Required** - Fast inference approach\n\n---\n\nThis notebook implements a **TBM pipeline** based on these proven strategies.","metadata":{}},{"cell_type":"markdown","source":"## Table of Contents\n\n1. [Setup & Installation](#setup)\n2. [Data Loading & Exploration](#data)\n3. [Composite Similarity Search](#search)\n4. [BioPython Sequence Alignment](#alignment)\n5. [Coordinate Transfer & Gap Filling](#transfer)\n6. [Adaptive Refinement](#refinement)\n7. [Prediction Pipeline](#pipeline)\n8. [Validation & TM-Score](#validation)\n9. [Submission Generation](#submission)","metadata":{}},{"cell_type":"markdown","source":"<a id=\"setup\"></a>\n## 1. Setup & Installation","metadata":{}},{"cell_type":"code","source":"# Install dependencies\n!pip install -q biopython pandas numpy scipy tqdm plotly scikit-learn","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T16:09:21.281468Z","iopub.execute_input":"2026-01-09T16:09:21.281827Z","iopub.status.idle":"2026-01-09T16:10:30.105106Z","shell.execute_reply.started":"2026-01-09T16:09:21.281796Z","shell.execute_reply":"2026-01-09T16:10:30.103621Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Core imports\nimport os\nimport gc\nimport warnings\nfrom pathlib import Path\nfrom typing import Dict, List, Tuple, Optional, Set\nfrom dataclasses import dataclass, field\nfrom collections import defaultdict\nfrom functools import lru_cache\n\nimport numpy as np\nimport pandas as pd\nfrom scipy.spatial.distance import cdist\nfrom scipy.spatial.transform import Rotation\nfrom scipy.optimize import minimize\nfrom sklearn.decomposition import PCA\nfrom tqdm.auto import tqdm\n\n# BioPython for alignment (as used by 1st place winner)\nfrom Bio import pairwise2\nfrom Bio.pairwise2 import format_alignment\nfrom Bio.Align import substitution_matrices\nfrom Bio.PDB import MMCIFParser\nfrom Bio import SeqIO\n\n# Visualization\nimport plotly.graph_objects as go\nimport plotly.express as px\nfrom plotly.subplots import make_subplots\n\nwarnings.filterwarnings('ignore')\n\n@dataclass\nclass Config:\n    \"\"\"Configuration for TBM pipeline.\"\"\"\n    # Paths\n    DATA_DIR: str = '/kaggle/input/stanford-rna-3d-folding-2'\n    OUTPUT_DIR: str = '/kaggle/working'\n    PDB_DIR: str = '/kaggle/input/stanford-rna-3d-folding-2/PDB_RNA'\n    \n    # TBM Parameters (from Part 1 winning approach)\n    NUM_PREDICTIONS: int = 5\n    MIN_IDENTITY: float = 0.25\n    MAX_TEMPLATES: int = 20\n    \n    # Composite Similarity Weights (from 1st place)\n    WEIGHT_GLOBAL: float = 0.4\n    WEIGHT_LOCAL: float = 0.3\n    WEIGHT_FEATURE: float = 0.2\n    WEIGHT_KMER: float = 0.1\n    \n    # Structure parameters\n    C1_C1_DISTANCE: float = 5.9  # Average C1'-C1' distance in Angstroms\n    GAP_PENALTY: float = -10.0\n    EXTEND_PENALTY: float = -0.5\n    \n    # Refinement\n    CLASH_THRESHOLD: float = 3.0  # Minimum distance before clash\n    REFINEMENT_ITERATIONS: int = 100\n\ncfg = Config()\nos.makedirs(cfg.OUTPUT_DIR, exist_ok=True)\n\nprint(\"✓ Setup complete!\")\nprint(f\"\\nComposite Similarity Weights:\")\nprint(f\"  Global:  {cfg.WEIGHT_GLOBAL:.0%}\")\nprint(f\"  Local:   {cfg.WEIGHT_LOCAL:.0%}\")\nprint(f\"  Feature: {cfg.WEIGHT_FEATURE:.0%}\")\nprint(f\"  K-mer:   {cfg.WEIGHT_KMER:.0%}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T16:10:30.10775Z","iopub.execute_input":"2026-01-09T16:10:30.108114Z","iopub.status.idle":"2026-01-09T16:10:36.074424Z","shell.execute_reply.started":"2026-01-09T16:10:30.108075Z","shell.execute_reply":"2026-01-09T16:10:36.073225Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"data\"></a>\n## 2. Data Loading & Exploration","metadata":{}},{"cell_type":"code","source":"# Load competition data\ntrain_sequences = pd.read_csv(f'{cfg.DATA_DIR}/train_sequences.csv')\ntrain_labels = pd.read_csv(f'{cfg.DATA_DIR}/train_labels.csv')\nval_sequences = pd.read_csv(f'{cfg.DATA_DIR}/validation_sequences.csv')\nval_labels = pd.read_csv(f'{cfg.DATA_DIR}/validation_labels.csv')\ntest_sequences = pd.read_csv(f'{cfg.DATA_DIR}/test_sequences.csv')\nsample_submission = pd.read_csv(f'{cfg.DATA_DIR}/sample_submission.csv')\n\nprint(\"=\"*60)\nprint(\"DATASET SUMMARY\")\nprint(\"=\"*60)\nprint(f\"Train sequences:      {len(train_sequences):,}\")\nprint(f\"Validation sequences: {len(val_sequences):,}\")\nprint(f\"Test sequences:       {len(test_sequences):,}\")\nprint(f\"Total templates:      {len(train_sequences) + len(val_sequences):,}\")\n\n# Sequence length statistics\nall_lens = pd.concat([\n    train_sequences['sequence'].str.len(),\n    val_sequences['sequence'].str.len(),\n    test_sequences['sequence'].str.len()\n])\nprint(f\"\\nSequence lengths: {all_lens.min()} - {all_lens.max()} (median: {all_lens.median():.0f})\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T16:10:36.076265Z","iopub.execute_input":"2026-01-09T16:10:36.077287Z","iopub.status.idle":"2026-01-09T16:10:47.280013Z","shell.execute_reply.started":"2026-01-09T16:10:36.077248Z","shell.execute_reply":"2026-01-09T16:10:47.278736Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Build comprehensive template library\n# Part 1 insight: Larger template library = better results\n\nclass TemplateLibrary:\n    \"\"\"Template library for TBM, combining train and validation data.\"\"\"\n    \n    def __init__(self):\n        self.sequences = {}  # target_id -> sequence\n        self.coords = {}     # target_id -> (L, 3) coordinates\n        self.features = {}   # target_id -> feature dict\n        \n    def add_from_dataframes(self, seq_df: pd.DataFrame, labels_df: pd.DataFrame):\n        \"\"\"Add templates from sequence and labels dataframes.\"\"\"\n        # Index labels by target\n        labels_df = labels_df.copy()\n        labels_df['target_id'] = labels_df['ID'].str.rsplit('_', n=1).str[0]\n        \n        for _, row in tqdm(seq_df.iterrows(), total=len(seq_df), desc='Loading templates'):\n            target_id = row['target_id']\n            sequence = row['sequence']\n            \n            # Get coordinates\n            target_labels = labels_df[labels_df['target_id'] == target_id].sort_values('resid')\n            \n            if len(target_labels) != len(sequence):\n                continue\n                \n            coords = np.stack([\n                target_labels['x_1'].values,\n                target_labels['y_1'].values,\n                target_labels['z_1'].values\n            ], axis=-1)\n            \n            # Skip if any NaN coordinates\n            if np.isnan(coords).any():\n                continue\n            \n            self.sequences[target_id] = sequence\n            self.coords[target_id] = coords\n            \n            # Precompute features\n            self.features[target_id] = self._compute_features(sequence)\n    \n    def _compute_features(self, sequence: str) -> Dict:\n        \"\"\"Compute sequence features for similarity scoring.\"\"\"\n        n = len(sequence)\n        \n        # Base composition\n        gc_content = (sequence.count('G') + sequence.count('C')) / n\n        au_content = (sequence.count('A') + sequence.count('U')) / n\n        \n        # Dinucleotide frequencies\n        dinucs = defaultdict(int)\n        for i in range(n - 1):\n            dinucs[sequence[i:i+2]] += 1\n        \n        # K-mers (3-mers)\n        kmers = set()\n        for i in range(n - 2):\n            kmers.add(sequence[i:i+3])\n        \n        return {\n            'gc_content': gc_content,\n            'au_content': au_content,\n            'length': n,\n            'dinucs': dict(dinucs),\n            'kmers': kmers\n        }\n    \n    def __len__(self):\n        return len(self.sequences)\n\n# Build template library (combining train + validation for maximum coverage)\nprint(\"\\nBuilding template library...\")\ntemplate_lib = TemplateLibrary()\ntemplate_lib.add_from_dataframes(train_sequences, train_labels)\ntemplate_lib.add_from_dataframes(val_sequences, val_labels)\n\nprint(f\"\\n✓ Template library: {len(template_lib):,} templates\")\nprint(f\"  (Part 1 insight: 18k templates >> 5k templates in performance)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T16:10:47.282166Z","iopub.execute_input":"2026-01-09T16:10:47.282562Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"search\"></a>\n## 3. Composite Similarity Search\n\nImplementation of the 1st place composite similarity scoring:\n- **40% Global**: Overall sequence identity\n- **30% Local**: Best local alignment score\n- **20% Feature**: Sequence feature similarity (GC content, length)\n- **10% K-mer**: Shared 3-mer Jaccard similarity","metadata":{}},{"cell_type":"code","source":"class CompositeSimilaritySearcher:\n    \"\"\"Template search using composite similarity scoring (1st place approach).\"\"\"\n    \n    def __init__(\n        self, \n        template_lib: TemplateLibrary,\n        weight_global: float = 0.4,\n        weight_local: float = 0.3,\n        weight_feature: float = 0.2,\n        weight_kmer: float = 0.1\n    ):\n        self.template_lib = template_lib\n        self.weights = {\n            'global': weight_global,\n            'local': weight_local,\n            'feature': weight_feature,\n            'kmer': weight_kmer\n        }\n        \n    def global_identity(self, seq1: str, seq2: str) -> float:\n        \"\"\"Calculate global sequence identity.\"\"\"\n        # Quick ungapped alignment\n        min_len = min(len(seq1), len(seq2))\n        max_len = max(len(seq1), len(seq2))\n        \n        if max_len == 0:\n            return 0.0\n        \n        # Try different offsets and take best\n        best_matches = 0\n        for offset in range(-min(10, min_len//2), min(10, min_len//2) + 1):\n            if offset >= 0:\n                s1, s2 = seq1, seq2[offset:]\n            else:\n                s1, s2 = seq1[-offset:], seq2\n            \n            matches = sum(1 for a, b in zip(s1, s2) if a == b)\n            best_matches = max(best_matches, matches)\n        \n        return best_matches / max_len\n    \n    def local_similarity(self, seq1: str, seq2: str) -> float:\n        \"\"\"Calculate local alignment score (normalized).\"\"\"\n        # Use BioPython's local alignment (as used by 1st place winner)\n        # For speed, only align if sequences are similar in length\n        len_ratio = min(len(seq1), len(seq2)) / max(len(seq1), len(seq2))\n        if len_ratio < 0.3:\n            return 0.0\n        \n        # Simple scoring: match=2, mismatch=-1\n        alignments = pairwise2.align.localms(\n            seq1, seq2,\n            2, -1,  # match, mismatch\n            -0.5, -0.1,  # gap open, gap extend\n            one_alignment_only=True\n        )\n        \n        if not alignments:\n            return 0.0\n        \n        # Normalize by maximum possible score\n        max_score = 2 * min(len(seq1), len(seq2))\n        return alignments[0].score / max_score if max_score > 0 else 0.0\n    \n    def feature_similarity(self, feat1: Dict, feat2: Dict) -> float:\n        \"\"\"Calculate feature-based similarity.\"\"\"\n        # GC content similarity\n        gc_sim = 1 - abs(feat1['gc_content'] - feat2['gc_content'])\n        \n        # Length similarity (penalize large differences)\n        len_ratio = min(feat1['length'], feat2['length']) / max(feat1['length'], feat2['length'])\n        \n        # Dinucleotide similarity\n        all_dinucs = set(feat1['dinucs'].keys()) | set(feat2['dinucs'].keys())\n        if all_dinucs:\n            dinuc_sim = sum(\n                min(feat1['dinucs'].get(d, 0), feat2['dinucs'].get(d, 0))\n                for d in all_dinucs\n            ) / max(\n                sum(feat1['dinucs'].values()),\n                sum(feat2['dinucs'].values())\n            )\n        else:\n            dinuc_sim = 0.0\n        \n        return (gc_sim + len_ratio + dinuc_sim) / 3\n    \n    def kmer_similarity(self, kmers1: Set, kmers2: Set) -> float:\n        \"\"\"Calculate k-mer Jaccard similarity.\"\"\"\n        if not kmers1 or not kmers2:\n            return 0.0\n        \n        intersection = len(kmers1 & kmers2)\n        union = len(kmers1 | kmers2)\n        \n        return intersection / union if union > 0 else 0.0\n    \n    def composite_score(\n        self, \n        query_seq: str, \n        query_features: Dict,\n        template_id: str\n    ) -> Tuple[float, Dict]:\n        \"\"\"Calculate composite similarity score.\"\"\"\n        template_seq = self.template_lib.sequences[template_id]\n        template_features = self.template_lib.features[template_id]\n        \n        # Calculate component scores\n        scores = {\n            'global': self.global_identity(query_seq, template_seq),\n            'local': self.local_similarity(query_seq, template_seq),\n            'feature': self.feature_similarity(query_features, template_features),\n            'kmer': self.kmer_similarity(query_features['kmers'], template_features['kmers'])\n        }\n        \n        # Weighted composite\n        composite = sum(\n            self.weights[k] * scores[k] \n            for k in scores\n        )\n        \n        return composite, scores\n    \n    def search(\n        self, \n        query_seq: str, \n        max_templates: int = 20,\n        min_score: float = 0.2\n    ) -> List[Tuple[str, float, Dict]]:\n        \"\"\"Search for similar templates.\n        \n        Returns:\n            List of (template_id, composite_score, component_scores)\n        \"\"\"\n        # Compute query features\n        query_features = self.template_lib._compute_features(query_seq)\n        \n        # Score all templates\n        results = []\n        for template_id in self.template_lib.sequences:\n            composite, scores = self.composite_score(\n                query_seq, query_features, template_id\n            )\n            if composite >= min_score:\n                results.append((template_id, composite, scores))\n        \n        # Sort by composite score\n        results.sort(key=lambda x: x[1], reverse=True)\n        \n        return results[:max_templates]\n\n# Initialize searcher\nsearcher = CompositeSimilaritySearcher(\n    template_lib,\n    weight_global=cfg.WEIGHT_GLOBAL,\n    weight_local=cfg.WEIGHT_LOCAL,\n    weight_feature=cfg.WEIGHT_FEATURE,\n    weight_kmer=cfg.WEIGHT_KMER\n)\n\nprint(\"✓ Composite similarity searcher initialized\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Test the searcher\ntest_query = test_sequences['sequence'].iloc[0]\nprint(f\"Query sequence length: {len(test_query)}\")\nprint(f\"Query: {test_query[:50]}...\")\n\ntemplates = searcher.search(test_query, max_templates=5, min_score=0.15)\n\nprint(f\"\\nTop {len(templates)} templates:\")\nprint(\"-\" * 70)\nprint(f\"{'Template':<15} {'Composite':>10} {'Global':>8} {'Local':>8} {'Feature':>8} {'Kmer':>8}\")\nprint(\"-\" * 70)\n\nfor template_id, composite, scores in templates:\n    print(f\"{template_id:<15} {composite:>10.3f} {scores['global']:>8.3f} {scores['local']:>8.3f} {scores['feature']:>8.3f} {scores['kmer']:>8.3f}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"alignment\"></a>\n## 4. BioPython Sequence Alignment\n\nThe 1st place winner used BioPython's aligner for accurate sequence alignment.","metadata":{}},{"cell_type":"code","source":"class BioPythonAligner:\n    \"\"\"Sequence alignment using BioPython (as used by 1st place winner).\"\"\"\n    \n    def __init__(\n        self, \n        match_score: float = 2.0,\n        mismatch_score: float = -1.0,\n        gap_open: float = -10.0,\n        gap_extend: float = -0.5\n    ):\n        self.match_score = match_score\n        self.mismatch_score = mismatch_score\n        self.gap_open = gap_open\n        self.gap_extend = gap_extend\n    \n    def align(self, query: str, template: str) -> Tuple[str, str, float]:\n        \"\"\"Perform global alignment.\n        \n        Returns:\n            (aligned_query, aligned_template, score)\n        \"\"\"\n        alignments = pairwise2.align.globalms(\n            query, template,\n            self.match_score, self.mismatch_score,\n            self.gap_open, self.gap_extend,\n            one_alignment_only=True\n        )\n        \n        if not alignments:\n            return query, template, 0.0\n        \n        best = alignments[0]\n        return best.seqA, best.seqB, best.score\n    \n    def get_alignment_mapping(self, query: str, template: str) -> Tuple[List[int], List[int]]:\n        \"\"\"Get position mapping from alignment.\n        \n        Returns:\n            (query_indices, template_indices) where -1 indicates a gap\n        \"\"\"\n        aligned_query, aligned_template, _ = self.align(query, template)\n        \n        query_indices = []\n        template_indices = []\n        \n        q_idx = 0\n        t_idx = 0\n        \n        for q_char, t_char in zip(aligned_query, aligned_template):\n            if q_char != '-' and t_char != '-':\n                # Match or mismatch\n                query_indices.append(q_idx)\n                template_indices.append(t_idx)\n                q_idx += 1\n                t_idx += 1\n            elif q_char == '-':\n                # Gap in query\n                t_idx += 1\n            else:\n                # Gap in template\n                query_indices.append(q_idx)\n                template_indices.append(-1)  # No template position\n                q_idx += 1\n        \n        # Handle any remaining query positions\n        while q_idx < len(query):\n            query_indices.append(q_idx)\n            template_indices.append(-1)\n            q_idx += 1\n        \n        return query_indices, template_indices\n\n# Initialize aligner\naligner = BioPythonAligner(\n    gap_open=cfg.GAP_PENALTY,\n    gap_extend=cfg.EXTEND_PENALTY\n)\n\nprint(\"✓ BioPython aligner initialized\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Test alignment\nif len(templates) > 0:\n    template_id = templates[0][0]\n    template_seq = template_lib.sequences[template_id]\n    \n    aligned_q, aligned_t, score = aligner.align(test_query[:100], template_seq[:100])\n    \n    print(f\"Alignment score: {score:.1f}\")\n    print(f\"\\nAligned query:    {aligned_q[:60]}...\")\n    print(f\"Aligned template: {aligned_t[:60]}...\")\n    \n    # Get mapping\n    q_indices, t_indices = aligner.get_alignment_mapping(test_query[:100], template_seq[:100])\n    mapped_count = sum(1 for t in t_indices if t >= 0)\n    print(f\"\\nMapped positions: {mapped_count}/{len(q_indices)} ({mapped_count/len(q_indices):.1%})\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"transfer\"></a>\n## 5. Coordinate Transfer & Gap Filling\n\nKey insight from Part 1: C1'-C1' distance is approximately **5.9 Angstroms** for gap filling.","metadata":{}},{"cell_type":"code","source":"class CoordinateTransfer:\n    \"\"\"Transfer and fill coordinates from template to query.\"\"\"\n    \n    def __init__(self, c1_c1_distance: float = 5.9):\n        self.c1_c1_distance = c1_c1_distance\n    \n    def transfer_coords(\n        self, \n        query_seq: str,\n        template_seq: str, \n        template_coords: np.ndarray,\n        aligner: BioPythonAligner\n    ) -> np.ndarray:\n        \"\"\"Transfer template coordinates to query using alignment.\n        \n        Returns:\n            (L, 3) array with NaN for unmapped positions\n        \"\"\"\n        query_len = len(query_seq)\n        coords = np.full((query_len, 3), np.nan)\n        \n        # Get alignment mapping\n        q_indices, t_indices = aligner.get_alignment_mapping(query_seq, template_seq)\n        \n        # Transfer coordinates\n        for q_idx, t_idx in zip(q_indices, t_indices):\n            if t_idx >= 0 and t_idx < len(template_coords):\n                if not np.isnan(template_coords[t_idx]).any():\n                    coords[q_idx] = template_coords[t_idx]\n        \n        return coords\n    \n    def fill_gaps(self, coords: np.ndarray) -> np.ndarray:\n        \"\"\"Fill gaps using C1'-C1' distance constraint.\n        \n        Part 1 insight: C1'-C1' distance ~5.9 Angstroms\n        \"\"\"\n        result = coords.copy()\n        n = len(coords)\n        \n        # Find valid positions\n        valid_mask = ~np.isnan(coords[:, 0])\n        valid_indices = np.where(valid_mask)[0]\n        \n        if len(valid_indices) == 0:\n            # No valid coordinates - generate de novo\n            return self._generate_helix(n)\n        \n        if len(valid_indices) == n:\n            # All coordinates valid\n            return result\n        \n        # Fill gaps using linear interpolation with distance constraint\n        for i in range(n):\n            if valid_mask[i]:\n                continue\n            \n            # Find nearest valid positions\n            prev_valid = None\n            next_valid = None\n            \n            for j in range(i - 1, -1, -1):\n                if valid_mask[j]:\n                    prev_valid = j\n                    break\n            \n            for j in range(i + 1, n):\n                if valid_mask[j]:\n                    next_valid = j\n                    break\n            \n            if prev_valid is not None and next_valid is not None:\n                # Interpolate between valid positions\n                gap_size = next_valid - prev_valid\n                t = (i - prev_valid) / gap_size\n                \n                # Linear interpolation\n                result[i] = (1 - t) * result[prev_valid] + t * result[next_valid]\n                \n                # Adjust to maintain C1'-C1' distance\n                if i > 0 and not np.isnan(result[i - 1]).any():\n                    direction = result[i] - result[i - 1]\n                    dist = np.linalg.norm(direction)\n                    if dist > 0:\n                        result[i] = result[i - 1] + direction / dist * self.c1_c1_distance\n                        \n            elif prev_valid is not None:\n                # Extend from previous position\n                direction = self._estimate_direction(result, prev_valid, 'forward')\n                result[i] = result[prev_valid] + direction * (i - prev_valid) * self.c1_c1_distance\n                \n            elif next_valid is not None:\n                # Extend from next position\n                direction = self._estimate_direction(result, next_valid, 'backward')\n                result[i] = result[next_valid] + direction * (next_valid - i) * self.c1_c1_distance\n        \n        return result\n    \n    def _estimate_direction(self, coords: np.ndarray, idx: int, direction: str) -> np.ndarray:\n        \"\"\"Estimate chain direction at a position.\"\"\"\n        valid_mask = ~np.isnan(coords[:, 0])\n        \n        if direction == 'forward':\n            # Look for previous valid positions\n            for i in range(idx - 1, max(0, idx - 5), -1):\n                if valid_mask[i]:\n                    vec = coords[idx] - coords[i]\n                    return vec / np.linalg.norm(vec)\n        else:\n            # Look for next valid positions\n            for i in range(idx + 1, min(len(coords), idx + 5)):\n                if valid_mask[i]:\n                    vec = coords[i] - coords[idx]\n                    return -vec / np.linalg.norm(vec)\n        \n        # Default direction\n        return np.array([1.0, 0.0, 0.0])\n    \n    def _generate_helix(self, n: int) -> np.ndarray:\n        \"\"\"Generate A-form RNA helix coordinates.\"\"\"\n        coords = np.zeros((n, 3))\n        radius = 10.5  # C1' radius\n        rise = 2.8  # Rise per base\n        twist = np.radians(32.7)  # Twist per base\n        \n        for i in range(n):\n            angle = twist * i\n            coords[i] = [\n                radius * np.cos(angle),\n                radius * np.sin(angle),\n                rise * i\n            ]\n        \n        # Center the structure\n        coords -= coords.mean(axis=0)\n        \n        return coords\n\n# Initialize coordinate transfer\ncoord_transfer = CoordinateTransfer(c1_c1_distance=cfg.C1_C1_DISTANCE)\n\nprint(\"✓ Coordinate transfer initialized\")\nprint(f\"  C1'-C1' distance: {cfg.C1_C1_DISTANCE} Å\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Test coordinate transfer\nif len(templates) > 0:\n    template_id = templates[0][0]\n    template_seq = template_lib.sequences[template_id]\n    template_coords = template_lib.coords[template_id]\n    \n    # Transfer coordinates\n    transferred = coord_transfer.transfer_coords(\n        test_query, template_seq, template_coords, aligner\n    )\n    \n    valid_before = np.sum(~np.isnan(transferred[:, 0]))\n    print(f\"Before gap filling: {valid_before}/{len(test_query)} positions ({valid_before/len(test_query):.1%})\")\n    \n    # Fill gaps\n    filled = coord_transfer.fill_gaps(transferred)\n    valid_after = np.sum(~np.isnan(filled[:, 0]))\n    print(f\"After gap filling:  {valid_after}/{len(test_query)} positions ({valid_after/len(test_query):.1%})\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"refinement\"></a>\n## 6. Adaptive Refinement\n\nLocal optimization to reduce steric clashes (from 1st place approach).","metadata":{}},{"cell_type":"code","source":"class AdaptiveRefinement:\n    \"\"\"Adaptive refinement to optimize structure quality.\"\"\"\n    \n    def __init__(\n        self, \n        clash_threshold: float = 3.0,\n        bond_distance: float = 5.9,\n        max_iterations: int = 100\n    ):\n        self.clash_threshold = clash_threshold\n        self.bond_distance = bond_distance\n        self.max_iterations = max_iterations\n    \n    def count_clashes(self, coords: np.ndarray) -> int:\n        \"\"\"Count number of steric clashes.\"\"\"\n        n = len(coords)\n        clashes = 0\n        \n        for i in range(n):\n            for j in range(i + 3, n):  # Skip adjacent positions\n                dist = np.linalg.norm(coords[i] - coords[j])\n                if dist < self.clash_threshold:\n                    clashes += 1\n        \n        return clashes\n    \n    def compute_energy(self, coords: np.ndarray) -> float:\n        \"\"\"Compute pseudo-energy for optimization.\"\"\"\n        n = len(coords)\n        energy = 0.0\n        \n        # Bond length penalty\n        for i in range(n - 1):\n            dist = np.linalg.norm(coords[i + 1] - coords[i])\n            energy += (dist - self.bond_distance) ** 2\n        \n        # Clash penalty\n        for i in range(n):\n            for j in range(i + 3, n):\n                dist = np.linalg.norm(coords[i] - coords[j])\n                if dist < self.clash_threshold:\n                    energy += (self.clash_threshold - dist) ** 2 * 10\n        \n        return energy\n    \n    def refine(self, coords: np.ndarray, fixed_mask: Optional[np.ndarray] = None) -> np.ndarray:\n        \"\"\"Refine coordinates to reduce clashes and improve geometry.\"\"\"\n        result = coords.copy()\n        n = len(coords)\n        \n        if fixed_mask is None:\n            fixed_mask = np.zeros(n, dtype=bool)\n        \n        initial_clashes = self.count_clashes(coords)\n        \n        # Simple gradient descent refinement\n        step_size = 0.1\n        \n        for iteration in range(self.max_iterations):\n            forces = np.zeros_like(result)\n            \n            # Bond length forces\n            for i in range(n - 1):\n                vec = result[i + 1] - result[i]\n                dist = np.linalg.norm(vec)\n                if dist > 0:\n                    force = (dist - self.bond_distance) * vec / dist\n                    if not fixed_mask[i]:\n                        forces[i] += force\n                    if not fixed_mask[i + 1]:\n                        forces[i + 1] -= force\n            \n            # Clash repulsion forces\n            for i in range(n):\n                for j in range(i + 3, n):\n                    vec = result[j] - result[i]\n                    dist = np.linalg.norm(vec)\n                    if dist < self.clash_threshold and dist > 0:\n                        repulsion = (self.clash_threshold - dist) * vec / dist * 2\n                        if not fixed_mask[i]:\n                            forces[i] -= repulsion\n                        if not fixed_mask[j]:\n                            forces[j] += repulsion\n            \n            # Apply forces\n            max_force = np.max(np.abs(forces))\n            if max_force < 0.01:\n                break\n            \n            result += forces * step_size\n        \n        final_clashes = self.count_clashes(result)\n        \n        return result\n    \n    def superpose(self, mobile: np.ndarray, target: np.ndarray) -> np.ndarray:\n        \"\"\"Superpose mobile onto target using Kabsch algorithm.\"\"\"\n        # Center both structures\n        mobile_centered = mobile - mobile.mean(axis=0)\n        target_centered = target - target.mean(axis=0)\n        \n        # Compute optimal rotation\n        H = mobile_centered.T @ target_centered\n        U, S, Vt = np.linalg.svd(H)\n        \n        # Handle reflection\n        d = np.sign(np.linalg.det(Vt.T @ U.T))\n        D = np.diag([1, 1, d])\n        \n        R = Vt.T @ D @ U.T\n        \n        # Apply rotation and translation\n        result = mobile_centered @ R.T + target.mean(axis=0)\n        \n        return result\n\n# Initialize refinement\nrefinement = AdaptiveRefinement(\n    clash_threshold=cfg.CLASH_THRESHOLD,\n    bond_distance=cfg.C1_C1_DISTANCE,\n    max_iterations=cfg.REFINEMENT_ITERATIONS\n)\n\nprint(\"✓ Adaptive refinement initialized\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Test refinement\nif 'filled' in dir():\n    clashes_before = refinement.count_clashes(filled)\n    refined = refinement.refine(filled)\n    clashes_after = refinement.count_clashes(refined)\n    \n    print(f\"Clashes before refinement: {clashes_before}\")\n    print(f\"Clashes after refinement:  {clashes_after}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"pipeline\"></a>\n## 7. Prediction Pipeline\n\nComplete TBM pipeline combining all components.","metadata":{}},{"cell_type":"code","source":"class WinningTBMPredictor:\n    \"\"\"Complete TBM predictor based on Part 1 winning approach.\"\"\"\n    \n    def __init__(\n        self,\n        template_lib: TemplateLibrary,\n        searcher: CompositeSimilaritySearcher,\n        aligner: BioPythonAligner,\n        coord_transfer: CoordinateTransfer,\n        refinement: AdaptiveRefinement,\n        num_predictions: int = 5,\n        min_score: float = 0.2,\n        max_templates: int = 20\n    ):\n        self.template_lib = template_lib\n        self.searcher = searcher\n        self.aligner = aligner\n        self.coord_transfer = coord_transfer\n        self.refinement = refinement\n        self.num_predictions = num_predictions\n        self.min_score = min_score\n        self.max_templates = max_templates\n    \n    def predict(self, query_seq: str) -> np.ndarray:\n        \"\"\"Predict 3D coordinates for a query sequence.\n        \n        Returns:\n            (L, num_predictions, 3) array of coordinates\n        \"\"\"\n        seq_len = len(query_seq)\n        predictions = []\n        \n        # Step 1: Search for templates\n        templates = self.searcher.search(\n            query_seq,\n            max_templates=self.max_templates,\n            min_score=self.min_score\n        )\n        \n        # Step 2: Generate predictions from top templates\n        for template_id, score, _ in templates:\n            if len(predictions) >= self.num_predictions:\n                break\n            \n            template_seq = self.template_lib.sequences[template_id]\n            template_coords = self.template_lib.coords[template_id]\n            \n            # Transfer coordinates\n            transferred = self.coord_transfer.transfer_coords(\n                query_seq, template_seq, template_coords, self.aligner\n            )\n            \n            # Fill gaps\n            filled = self.coord_transfer.fill_gaps(transferred)\n            \n            # Skip if too many gaps\n            coverage = np.sum(~np.isnan(transferred[:, 0])) / seq_len\n            if coverage < 0.2:\n                continue\n            \n            # Refine\n            refined = self.refinement.refine(filled)\n            \n            predictions.append(refined)\n        \n        # Step 3: Generate diverse predictions if needed\n        while len(predictions) < self.num_predictions:\n            if len(predictions) > 0:\n                # Add noise to existing predictions\n                base = predictions[len(predictions) % len(predictions)]\n                noise_scale = 0.5 + 0.5 * len(predictions)\n                noisy = base + np.random.randn(*base.shape) * noise_scale\n                predictions.append(noisy)\n            else:\n                # Generate de novo helix\n                helix = self.coord_transfer._generate_helix(seq_len)\n                noise = np.random.randn(*helix.shape) * 0.5 * (len(predictions) + 1)\n                predictions.append(helix + noise)\n        \n        # Stack predictions: (L, num_predictions, 3)\n        result = np.stack(predictions[:self.num_predictions], axis=1)\n        \n        return result\n\n# Initialize predictor\npredictor = WinningTBMPredictor(\n    template_lib=template_lib,\n    searcher=searcher,\n    aligner=aligner,\n    coord_transfer=coord_transfer,\n    refinement=refinement,\n    num_predictions=cfg.NUM_PREDICTIONS,\n    min_score=cfg.MIN_IDENTITY,\n    max_templates=cfg.MAX_TEMPLATES\n)\n\nprint(\"✓ Winning TBM predictor initialized\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Test prediction on first test sequence\ntest_prediction = predictor.predict(test_query)\n\nprint(f\"Prediction shape: {test_prediction.shape}\")\nprint(f\"Expected: ({len(test_query)}, {cfg.NUM_PREDICTIONS}, 3)\")\n\n# Visualize\nfig = go.Figure()\n\ncolors = px.colors.qualitative.Set1\nfor i in range(cfg.NUM_PREDICTIONS):\n    coords = test_prediction[:, i, :]\n    fig.add_trace(go.Scatter3d(\n        x=coords[:, 0],\n        y=coords[:, 1],\n        z=coords[:, 2],\n        mode='markers+lines',\n        name=f'Prediction {i+1}',\n        marker=dict(size=3, color=colors[i % len(colors)]),\n        line=dict(width=2, color=colors[i % len(colors)])\n    ))\n\nfig.update_layout(\n    title=f'5 Predicted RNA Structures (TBM Approach)',\n    scene=dict(\n        xaxis_title='X (Å)',\n        yaxis_title='Y (Å)',\n        zaxis_title='Z (Å)',\n        aspectmode='data'\n    ),\n    width=900,\n    height=700\n)\nfig.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"validation\"></a>\n## 8. Validation & TM-Score\n\nValidate on validation set to estimate performance.","metadata":{}},{"cell_type":"code","source":"def compute_tm_score(pred_coords: np.ndarray, true_coords: np.ndarray) -> float:\n    \"\"\"Compute TM-score between predicted and true coordinates.\n    \n    TM-score is normalized by target length and is in [0, 1].\n    \"\"\"\n    L = len(true_coords)\n    \n    if L < 5:\n        return 0.0\n    \n    # Remove any NaN values\n    valid_mask = ~np.isnan(pred_coords[:, 0]) & ~np.isnan(true_coords[:, 0])\n    if valid_mask.sum() < 5:\n        return 0.0\n    \n    pred = pred_coords[valid_mask]\n    true = true_coords[valid_mask]\n    L_aligned = len(pred)\n    \n    # d0 parameter (length-dependent)\n    d0 = 1.24 * (L - 15) ** (1/3) - 1.8\n    d0 = max(d0, 0.5)\n    \n    # Superpose predicted onto true\n    pred_centered = pred - pred.mean(axis=0)\n    true_centered = true - true.mean(axis=0)\n    \n    # Kabsch algorithm\n    H = pred_centered.T @ true_centered\n    U, S, Vt = np.linalg.svd(H)\n    d = np.sign(np.linalg.det(Vt.T @ U.T))\n    D = np.diag([1, 1, d])\n    R = Vt.T @ D @ U.T\n    \n    pred_rotated = pred_centered @ R.T\n    \n    # Compute TM-score\n    distances = np.linalg.norm(pred_rotated - true_centered, axis=1)\n    tm_score = np.sum(1 / (1 + (distances / d0) ** 2)) / L\n    \n    return tm_score\n\ndef evaluate_predictions(\n    predictions: np.ndarray, \n    true_coords: np.ndarray\n) -> Tuple[float, int]:\n    \"\"\"Evaluate predictions and return best TM-score.\n    \n    Returns:\n        (best_tm_score, best_prediction_index)\n    \"\"\"\n    num_preds = predictions.shape[1]\n    best_score = 0.0\n    best_idx = 0\n    \n    for i in range(num_preds):\n        score = compute_tm_score(predictions[:, i, :], true_coords)\n        if score > best_score:\n            best_score = score\n            best_idx = i\n    \n    return best_score, best_idx\n\nprint(\"✓ TM-score functions defined\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Validate on a subset of validation data\n# ⚠️ IMPORTANT DISCLAIMER:\n# Since we use validation data AS TEMPLATES, validation scores will be \n# OPTIMISTICALLY BIASED (near-exact matches possible).\n# Real test set performance will likely be LOWER than these numbers.\n# This validation is for sanity checking, NOT performance estimation.\n\n# Build validation coordinate lookup\nval_coords = {}\nval_labels['target_id'] = val_labels['ID'].str.rsplit('_', n=1).str[0]\n\nfor target_id in val_sequences['target_id'].unique():\n    target_labels = val_labels[val_labels['target_id'] == target_id].sort_values('resid')\n    if len(target_labels) > 0:\n        coords = np.stack([\n            target_labels['x_1'].values,\n            target_labels['y_1'].values,\n            target_labels['z_1'].values\n        ], axis=-1)\n        val_coords[target_id] = coords\n\n# Evaluate on sample\nsample_size = min(20, len(val_sequences))\ntm_scores = []\n\nprint(f\"Evaluating on {sample_size} validation sequences...\")\nprint(\"⚠️ Note: Scores are optimistic due to template overlap!\")\nprint(\"-\" * 50)\n\nfor idx in tqdm(range(sample_size)):\n    row = val_sequences.iloc[idx]\n    target_id = row['target_id']\n    sequence = row['sequence']\n    \n    if target_id not in val_coords:\n        continue\n    \n    # Predict\n    pred = predictor.predict(sequence)\n    true = val_coords[target_id]\n    \n    if len(pred) != len(true):\n        continue\n    \n    # Evaluate\n    best_score, best_idx = evaluate_predictions(pred, true)\n    tm_scores.append(best_score)\n\nif tm_scores:\n    print(f\"\\nValidation Results (OPTIMISTIC - see disclaimer above):\")\n    print(f\"  Mean TM-score: {np.mean(tm_scores):.4f}\")\n    print(f\"  Median TM-score: {np.median(tm_scores):.4f}\")\n    print(f\"  Min TM-score: {np.min(tm_scores):.4f}\")\n    print(f\"  Max TM-score: {np.max(tm_scores):.4f}\")\n    print(f\"\\n⚠️ Expect LOWER scores on the actual test set!\")\nelse:\n    print(\"No valid validation samples\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"<a id=\"submission\"></a>\n## 9. Submission Generation","metadata":{}},{"cell_type":"code","source":"def generate_submission(\n    test_sequences: pd.DataFrame,\n    predictor: WinningTBMPredictor,\n    num_predictions: int = 5\n) -> pd.DataFrame:\n    \"\"\"Generate submission dataframe.\"\"\"\n    all_rows = []\n    \n    for idx in tqdm(range(len(test_sequences)), desc='Generating predictions'):\n        row = test_sequences.iloc[idx]\n        target_id = row['target_id']\n        sequence = row['sequence']\n        seq_len = len(sequence)\n        \n        # Get predictions\n        predictions = predictor.predict(sequence)\n        \n        # Create submission rows\n        for i in range(seq_len):\n            resid = i + 1\n            resname = sequence[i]\n            \n            pred_row = {\n                'ID': f'{target_id}_{resid}',\n                'resname': resname,\n                'resid': resid\n            }\n            \n            for j in range(num_predictions):\n                # Clip coordinates to valid range\n                x = np.clip(predictions[i, j, 0], -999.999, 9999.999)\n                y = np.clip(predictions[i, j, 1], -999.999, 9999.999)\n                z = np.clip(predictions[i, j, 2], -999.999, 9999.999)\n                \n                pred_row[f'x_{j+1}'] = round(x, 3)\n                pred_row[f'y_{j+1}'] = round(y, 3)\n                pred_row[f'z_{j+1}'] = round(z, 3)\n            \n            all_rows.append(pred_row)\n    \n    return pd.DataFrame(all_rows)\n\n# Generate submission\nprint(\"Generating submission...\")\nsubmission_df = generate_submission(test_sequences, predictor, cfg.NUM_PREDICTIONS)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Validate submission format\nprint(\"\\n\" + \"=\"*60)\nprint(\"SUBMISSION VALIDATION\")\nprint(\"=\"*60)\n\nexpected_cols = ['ID', 'resname', 'resid']\nfor i in range(1, 6):\n    expected_cols.extend([f'x_{i}', f'y_{i}', f'z_{i}'])\n\nmissing_cols = set(expected_cols) - set(submission_df.columns)\nif missing_cols:\n    print(f\"Missing columns: {missing_cols}\")\nelse:\n    print(\"✓ All required columns present\")\n\ncoord_cols = [c for c in submission_df.columns if c.startswith(('x_', 'y_', 'z_'))]\nprint(f\"\\nCoordinate statistics:\")\nprint(f\"  Min: {submission_df[coord_cols].min().min():.3f}\")\nprint(f\"  Max: {submission_df[coord_cols].max().max():.3f}\")\nprint(f\"  Mean: {submission_df[coord_cols].mean().mean():.3f}\")\n\nnan_count = submission_df[coord_cols].isna().sum().sum()\nprint(f\"  NaN values: {nan_count}\")\n\nprint(f\"\\nSubmission shape: {submission_df.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Reorder and save\nordered_cols = ['ID', 'resname', 'resid']\nfor i in range(1, 6):\n    ordered_cols.extend([f'x_{i}', f'y_{i}', f'z_{i}'])\n\nsubmission_df = submission_df[ordered_cols]\nsubmission_df.to_csv('submission.csv', index=False)\n\nprint(f\"\\n✓ Submission saved to 'submission.csv'\")\nprint(f\"  File size: {os.path.getsize('submission.csv') / 1024:.1f} KB\")\nprint(f\"  Total rows: {len(submission_df):,}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Final visualization\nsample_target = test_sequences['target_id'].iloc[0]\nsample_pred = submission_df[submission_df['ID'].str.startswith(sample_target + '_')]\n\nfig = make_subplots(\n    rows=1, cols=2,\n    specs=[[{'type': 'scatter3d'}, {'type': 'scatter3d'}]],\n    subplot_titles=('Prediction 1', 'Prediction 2')\n)\n\nfor idx, pred_num in enumerate([1, 2], 1):\n    x = sample_pred[f'x_{pred_num}'].values\n    y = sample_pred[f'y_{pred_num}'].values\n    z = sample_pred[f'z_{pred_num}'].values\n    \n    fig.add_trace(\n        go.Scatter3d(\n            x=x, y=y, z=z,\n            mode='markers+lines',\n            marker=dict(\n                size=4,\n                color=np.arange(len(x)),\n                colorscale='Viridis'\n            ),\n            line=dict(width=2, color='lightgray')\n        ),\n        row=1, col=idx\n    )\n\nfig.update_layout(\n    title=f'Sample Predictions for {sample_target}',\n    height=600,\n    showlegend=False\n)\nfig.show()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"---\n\n## Summary\n\nThis notebook implements a **TBM approach inspired by Part 1 winning strategies**:\n\n### Key Components\n\n1. **Composite Similarity Search** (40% global + 30% local + 20% feature + 10% kmer)\n2. **BioPython Alignment** (as used by 1st place winner)\n3. **Gap Filling** with C1'-C1' distance ~5.9Å\n4. **Adaptive Refinement** to reduce steric clashes\n\n### Part 1 Insights Applied\n\n| Part 1 Finding | Our Implementation |\n|----------------|-------------------|\n| TBM-only performed well | Pure TBM pipeline (no deep learning) |\n| Template library size matters | Combined train + validation templates |\n| Composite scoring > simple identity | 4-component weighted scoring |\n| BioPython aligner works well | Used for all alignments |\n| C1'-C1' ~5.9Å for gap filling | Implemented in coordinate transfer |\n\n### Important Notes\n\n⚠️ **Performance Disclaimer**:\n- Part 1 top solutions reported TM-scores around 0.35-0.59\n- **Your actual score depends heavily on template coverage** for test sequences\n- Sequences without good templates will score lower\n- The validation scores in this notebook are **optimistically biased** (template overlap)\n\n### Why TBM Works\n\n- Similar sequences often fold into similar structures\n- No GPU training required - fast inference\n- Leverages existing structural knowledge from PDB\n- Suitable for Kaggle's 8-hour time limit\n\n---\n\n**If you found this notebook helpful, please upvote!**\n\nGood luck in the competition!","metadata":{}}]}