{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"datasetVersion","sourceId":14604295,"datasetId":9328538,"databundleVersionId":15440074}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":2354.045049,"end_time":"2026-01-13T06:54:44.524542","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-01-13T06:15:30.479493","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install --no-index /kaggle/input/biopython-cp312/biopython-1.86-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-21T23:59:49.724626Z","iopub.execute_input":"2026-03-21T23:59:49.724872Z","iopub.status.idle":"2026-03-21T23:59:56.255202Z","shell.execute_reply.started":"2026-03-21T23:59:49.724849Z","shell.execute_reply":"2026-03-21T23:59:56.254222Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport random\nimport time\nimport warnings\nimport os, sys\nfrom tqdm import tqdm\n\nwarnings.filterwarnings('ignore')\n\n# =============================================================================\n# 1. DATA LOADING\n# =============================================================================\nDATA_PATH = '/kaggle/input/stanford-rna-3d-folding-2/'\ntrain_seqs   = pd.read_csv(DATA_PATH + 'train_sequences.csv')\ntest_seqs    = pd.read_csv(DATA_PATH + 'test_sequences.csv')\ntrain_labels = pd.read_csv(DATA_PATH + 'train_labels.csv')\n\nsys.path.append(os.path.join(DATA_PATH, \"extra\"))\n\n# =============================================================================\n# 2. FASTA / STOICHIOMETRY PARSING\n# =============================================================================\n\n# --- Robust import for Kaggle's extra/parse_fasta_py.py (it may miss typing imports) ---\ntry:\n    import typing as _typing\n    import builtins as _builtins\n\n    # Make these names available during module import-time annotation evaluation\n    _builtins.Dict  = getattr(_typing, \"Dict\")\n    _builtins.Tuple = getattr(_typing, \"Tuple\")\n    _builtins.List  = getattr(_typing, \"List\")\n\n    from parse_fasta_py import parse_fasta as _parse_fasta_raw\n\n    # Normalize output to: {chain_id: sequence_string}\n    def parse_fasta(fasta_content: str):\n        d = _parse_fasta_raw(fasta_content)\n        out = {}\n        for k, v in d.items():\n            # some variants return (sequence, headers/lines) or similar\n            out[k] = v[0] if isinstance(v, tuple) else v\n        return out\n\nexcept Exception:\n    # Fallback FASTA parser: {chain_id: sequence_string}\n    def parse_fasta(fasta_content: str):\n        out = {}\n        cur = None\n        seq_parts = []\n        for line in str(fasta_content).splitlines():\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                if cur is not None:\n                    out[cur] = \"\".join(seq_parts)\n                header = line[1:]\n                # First token is usually chain id in this dataset\n                cur = header.split()[0]\n                seq_parts = []\n            else:\n                seq_parts.append(line.replace(\" \", \"\"))\n        if cur is not None:\n            out[cur] = \"\".join(seq_parts)\n        return out\n\n\ndef parse_stoichiometry(stoich: str):\n    if pd.isna(stoich) or str(stoich).strip() == \"\":\n        return []\n    out = []\n    for part in str(stoich).split(';'):\n        ch, cnt = part.split(':')\n        out.append((ch.strip(), int(cnt)))\n    return out\n\n\ndef get_chain_segments(row):\n    \"\"\"\n    Returns list of (start, end) segments in row['sequence'] corresponding\n    to chain copies in stoichiometry order.\n    Falls back to a single segment if parsing fails.\n    \"\"\"\n    seq     = row['sequence']\n    stoich  = row.get('stoichiometry', '')\n    all_seq = row.get('all_sequences', '')\n\n    if pd.isna(stoich) or pd.isna(all_seq) or str(stoich).strip() == \"\" or str(all_seq).strip() == \"\":\n        return [(0, len(seq))]\n\n    try:\n        chain_dict = parse_fasta(all_seq)   # dict: chain_id -> sequence\n        order      = parse_stoichiometry(stoich)\n        segs       = []\n        pos        = 0\n        for ch, cnt in order:\n            base = chain_dict.get(ch)\n            if base is None:\n                return [(0, len(seq))]\n            for _ in range(cnt):\n                L = len(base)\n                segs.append((pos, pos + L))\n                pos += L\n        if pos != len(seq):\n            return [(0, len(seq))]\n        return segs\n    except Exception:\n        return [(0, len(seq))]\n\n\ndef build_segments_map(df):\n    seg_map    = {}\n    stoich_map = {}\n    for _, r in df.iterrows():\n        tid               = r['target_id']\n        seg_map[tid]      = get_chain_segments(r)\n        stoich_map[tid]   = str(r.get('stoichiometry', '') if not pd.isna(r.get('stoichiometry', '')) else '')\n    return seg_map, stoich_map\n\n\ntrain_segs_map, train_stoich_map = build_segments_map(train_seqs)\ntest_segs_map,  test_stoich_map  = build_segments_map(test_seqs)\n\n\ndef process_labels(labels_df):\n    coords_dict = {}\n    # Faster + safer prefix extraction\n    prefixes = labels_df['ID'].str.rsplit('_', n=1).str[0]\n    for id_prefix, group in labels_df.groupby(prefixes):\n        coords_dict[id_prefix] = group.sort_values('resid')[['x_1', 'y_1', 'z_1']].values\n    return coords_dict\n\n\ntrain_coords_dict = process_labels(train_labels)\n\n# =============================================================================\n# 3. SEQUENCE ALIGNER SETUP\n# =============================================================================\nfrom Bio.Align import PairwiseAligner\n\naligner      = PairwiseAligner()\naligner.mode = 'global'\n\n# --- Alignment scoring parameters ---\n# match_score=2      : Reward for matching bases; RNA is relatively conserved so keep it high\n# mismatch_score=-2  : Increased mismatch penalty; RNA structure is more sensitive to mismatches\n# open_gap_score=-10 : High gap-open penalty; gaps cause residue misalignment in coordinate mapping\n# extend_gap_score   : Small extension penalty to still allow necessary gaps\naligner.match_score      = 2\naligner.mismatch_score   = -2     # Changed from -1.5 → stricter mismatch penalty\naligner.open_gap_score   = -10    # Changed from -8  → reduce unwanted gaps\naligner.extend_gap_score = -0.5   # Changed from -0.4\n\n# Penalize terminal gaps to prevent end-gap semi-global behavior.\n# Ensures residue numbering stays aligned at sequence ends.\naligner.query_left_open_gap_score    = -8\naligner.query_left_extend_gap_score  = -0.4\naligner.query_right_open_gap_score   = -8\naligner.query_right_extend_gap_score = -0.4\naligner.target_left_open_gap_score   = -8\naligner.target_left_extend_gap_score = -0.4\naligner.target_right_open_gap_score  = -8\naligner.target_right_extend_gap_score = -0.4\n\n# Length similarity filter threshold.\n# Skip any training sequence whose length differs from the query by more than this ratio.\n# Original author used 0.3 (30%). Reduced to 0.25 for stricter template quality.\nLENGTH_DIFF_THRESHOLD = 0.25\n\n# =============================================================================\n# 4. CORE FUNCTIONS\n# =============================================================================\n\ndef find_similar_sequences(query_seq, train_seqs_df, train_coords_dict, top_n=5):\n    \"\"\"\n    Search the training set for sequences most similar to query_seq.\n\n    Steps:\n      1. Pre-filter by length (fast, avoids expensive alignment on hopeless candidates)\n      2. Compute fast alignment score (no traceback)\n      3. Normalize score by sequence length\n      4. Return top_n hits sorted by descending similarity\n    \"\"\"\n    similar_seqs = []\n\n    for _, row in train_seqs_df.iterrows():\n        target_id, train_seq = row['target_id'], row['sequence']\n        if target_id not in train_coords_dict:\n            continue\n\n        # Length pre-filter: skip sequences too different in length\n        if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > LENGTH_DIFF_THRESHOLD:\n            continue\n\n        # Fast score: calculates alignment score without full traceback overhead\n        raw_score        = aligner.score(query_seq, train_seq)\n        normalized_score = raw_score / (2 * min(len(query_seq), len(train_seq)))\n        similar_seqs.append((target_id, train_seq, normalized_score, train_coords_dict[target_id]))\n\n    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords):\n    \"\"\"\n    Map template 3D coordinates onto the query sequence using pairwise alignment.\n\n    Uses vectorized segment mapping (faster than character-by-character loop).\n    Gaps in the alignment are filled by linear interpolation between known\n    neighboring residues, or by linear extrapolation at the ends.\n    \"\"\"\n    # aligner.align returns an iterator; take only the first (best) alignment\n    alignment  = next(iter(aligner.align(query_seq, template_seq)))\n    new_coords = np.full((len(query_seq), 3), np.nan)\n\n    # alignment.aligned gives (start, end) segment pairs for matched regions.\n    # Vectorized slice assignment avoids slow Python character loops.\n    for (q_start, q_end), (t_start, t_end) in zip(*alignment.aligned):\n        t_chunk = template_coords[t_start:t_end]\n        if len(t_chunk) == (q_end - q_start):   # safety check\n            new_coords[q_start:q_end] = t_chunk\n\n    # --- Fill gaps by interpolation / extrapolation ---\n    for i in range(len(new_coords)):\n        if np.isnan(new_coords[i, 0]):\n            prev_v = next((j for j in range(i - 1, -1, -1)           if not np.isnan(new_coords[j, 0])), -1)\n            next_v = next((j for j in range(i + 1, len(new_coords))  if not np.isnan(new_coords[j, 0])), -1)\n\n            if prev_v >= 0 and next_v >= 0:\n                # Linear interpolation between two known neighbors\n                w              = (i - prev_v) / (next_v - prev_v)\n                new_coords[i]  = (1 - w) * new_coords[prev_v] + w * new_coords[next_v]\n            elif prev_v >= 0:\n                new_coords[i] = new_coords[prev_v] + [3, 0, 0]   # extrapolate forward\n            elif next_v >= 0:\n                new_coords[i] = new_coords[next_v] + [3, 0, 0]   # extrapolate backward\n            else:\n                new_coords[i] = [i * 3, 0, 0]                    # fallback: straight line\n\n    return np.nan_to_num(new_coords)\n\n\ndef adaptive_rna_constraints(coordinates, sequence, confidence=1.0):\n    \"\"\"\n    Apply soft physical constraints to refine predicted coordinates.\n\n    Applies two passes of bond-length correction:\n      - Nearest-neighbor (i, i+1): target distance ~5.95 Å  (P-to-P backbone)\n      - Next-nearest    (i, i+2): target distance ~10.2 Å\n    Correction strength scales with (1 - confidence):\n    high-confidence templates get minimal adjustment; low-confidence ones get more.\n    \"\"\"\n    refined_coords = coordinates.copy()\n    n              = len(sequence)\n    strength       = 0.68 * (1.0 - min(confidence, 0.96))\n\n    for _ in range(2):\n        for i in range(n - 1):\n            p1, p2 = refined_coords[i], refined_coords[i + 1]\n            dist   = np.linalg.norm(p2 - p1)\n            if dist > 0:\n                adj = (5.95 - dist) * strength * 0.45\n                refined_coords[i + 1] += (p2 - p1) / dist * adj\n\n            if i < n - 2:\n                p3    = refined_coords[i + 2]\n                dist2 = np.linalg.norm(p3 - p1)\n                if dist2 > 0:\n                    adj2 = (10.2 - dist2) * strength * 0.25\n                    refined_coords[i + 2] += (p3 - p1) / dist2 * adj2\n\n    return refined_coords\n\n\ndef add_position_aware_noise(coords, sim_score, pred_index, sequence):\n    \"\"\"\n    Apply position-aware noise perturbation to a coordinate prediction.\n\n    Motivation:\n      - RNA terminal regions (5' and 3' ends) are structurally more flexible\n        and less constrained than the core, so they receive amplified noise.\n      - Prediction 0 (best template) is kept clean with no noise.\n      - Subsequent predictions receive noise scaled by both similarity score\n        and residue position, increasing structural diversity in the ensemble.\n\n    Args:\n        coords      : (L, 3) coordinate array\n        sim_score   : normalized alignment similarity score (0 to 1)\n        pred_index  : slot index among the 5 predictions (0 = best/cleanest)\n        sequence    : RNA sequence string (length L)\n\n    Returns:\n        Perturbed coordinate array of shape (L, 3)\n    \"\"\"\n    if pred_index == 0:\n        return coords   # Best prediction stays completely clean\n\n    n = len(sequence)\n\n    # Base noise magnitude: lower similarity → larger perturbation\n    base_noise = max(0.01, (0.5 - sim_score) * 0.1)\n\n    # Position weights: terminal 20% of residues (5' and 3' ends) get 1.5× noise\n    # because RNA termini are intrinsically more flexible and less well-defined\n    position_weights              = np.ones(n)\n    tail_len                      = max(1, int(n * 0.2))\n    position_weights[:tail_len]   = 1.5   # 5' end\n    position_weights[-tail_len:]  = 1.5   # 3' end\n\n    noise  = np.random.normal(0, base_noise, coords.shape)\n    noise *= position_weights[:, np.newaxis]   # scale each residue independently\n    return coords + noise\n\n\ndef weighted_average_templates(similar_seqs, query_seq, top_k=3):\n    \"\"\"\n    Blend coordinate predictions from multiple templates using\n    similarity-squared weighting.\n\n    Motivation:\n      Relying on a single best template is fragile — a misaligned region\n      in the top match corrupts the whole prediction. Averaging across\n      multiple templates reduces this sensitivity and produces a more\n      robust consensus structure.\n\n      Similarity is squared before weighting so the best template still\n      dominates, but lower-ranked templates contribute proportionally.\n\n    Args:\n        similar_seqs : output list from find_similar_sequences()\n        query_seq    : RNA sequence string for the query\n        top_k        : number of top templates to blend (default: 3)\n\n    Returns:\n        Blended coordinate array of shape (L, 3)\n    \"\"\"\n    if len(similar_seqs) == 0:\n        # No templates found — return a simple linear chain as fallback\n        n      = len(query_seq)\n        coords = np.zeros((n, 3))\n        for j in range(1, n):\n            coords[j] = coords[j - 1] + [4.0, 0, 0]\n        return coords\n\n    use_k        = min(top_k, len(similar_seqs))\n    total_weight = 0.0\n    blended      = np.zeros((len(query_seq), 3))\n\n    for i in range(use_k):\n        t_id, t_seq, sim, t_coords = similar_seqs[i]\n        if sim <= 0:\n            continue\n        adapted       = adapt_template_to_query(query_seq, t_seq, t_coords)\n        weight         = sim ** 2    # squared weight: best template dominates\n        blended       += weight * adapted\n        total_weight  += weight\n\n    if total_weight > 0:\n        blended /= total_weight\n\n    return blended\n\n\ndef predict_rna_structures(sequence, target_id, train_seqs_df, train_coords_dict, n_predictions=5):\n    \"\"\"\n    Generate 5 structural predictions for a single RNA sequence.\n\n    Ensemble strategy:\n      - Prediction 0 : Similarity-weighted average of top-3 templates (most stable / robust)\n      - Prediction 1 : Best single template, clean (no noise)\n      - Predictions 2-4 : Lower-ranked templates + position-aware noise (structural diversity)\n\n    This is an improvement over the original code, which used one template per slot\n    with uniform noise and no multi-template blending.\n    \"\"\"\n    predictions  = []\n\n    # Fetch extra candidates so single-template slots still have choices\n    similar_seqs = find_similar_sequences(\n        sequence, train_seqs_df, train_coords_dict, top_n=n_predictions + 2\n    )\n\n    # --- Prediction 0: Multi-template weighted average ---\n    blended = weighted_average_templates(similar_seqs, sequence, top_k=3)\n    blended = adaptive_rna_constraints(blended, sequence, confidence=0.9)\n    predictions.append(blended)\n\n    # --- Predictions 1–4: Single template + position-aware noise ---\n    for i in range(1, n_predictions):\n        if i - 1 < len(similar_seqs):\n            t_id, t_seq, sim, t_coords = similar_seqs[i - 1]\n            adapted = adapt_template_to_query(sequence, t_seq, t_coords)\n            refined = adaptive_rna_constraints(adapted, sequence, confidence=sim)\n            refined = add_position_aware_noise(refined, sim, i, sequence)  # ← BUG FIX: was appending unperturbed coords\n            predictions.append(refined)\n        else:\n            # Fallback: straight-line chain when no template is available\n            n      = len(sequence)\n            coords = np.zeros((n, 3))\n            for j in range(1, n):\n                coords[j] = coords[j - 1] + [4.0, 0, 0]\n            predictions.append(coords)\n\n    return predictions\n\n# =============================================================================\n# 5. SUBMISSION VALIDATION\n# =============================================================================\n\ndef validate_submission(csv_path):\n    \"\"\"\n    Run sanity checks on the submission CSV before uploading to Kaggle.\n    Catches common issues that would cause a submission to score incorrectly.\n    \"\"\"\n    sub = pd.read_csv(csv_path)\n\n    print(\"=\" * 40)\n    print(\"   Submission Validation Report\")\n    print(\"=\" * 40)\n    print(f\"  Total rows : {len(sub)}\")\n    print(f\"  Columns    : {list(sub.columns)}\\n\")\n\n    checks = {\n        'No NaN values present'          : not sub.isnull().any().any(),\n        'Has all 5 coordinate sets (x/y/z)' : all(f'x_{i}' in sub.columns for i in range(1, 6)),\n        'resname only contains A/C/G/U'  : sub['resname'].str.upper().isin(['A', 'C', 'G', 'U']).all(),\n        'resid starts from 1'            : sub['resid'].min() == 1,\n        'Row count is greater than 0'    : len(sub) > 0,\n    }\n\n    all_pass = True\n    for check_name, result in checks.items():\n        icon = '✅' if result else '❌'\n        print(f\"  {icon}  {check_name}\")\n        if not result:\n            all_pass = False\n\n    print()\n    if all_pass:\n        print(\"  🎉 All checks passed — safe to submit!\")\n    else:\n        print(\"  ⚠️  Issues found — do NOT submit yet!\")\n    print(\"=\" * 40)\n\n    return all_pass\n\n# =============================================================================\n# 6. MAIN PREDICTION LOOP\n# =============================================================================\n\nall_predictions = []\nstart_time      = time.time()\nfailed_targets  = []   # Track any sequences that fail so we can inspect them\n\nfor idx, row in tqdm(test_seqs.iterrows(), total=len(test_seqs), desc=\"Predicting RNA structures\"):\n    tid, seq = row['target_id'], row['sequence']\n\n    try:\n        preds = predict_rna_structures(seq, tid, train_seqs, train_coords_dict)\n\n    except Exception as e:\n        print(f\"\\n❌ Failed on: {tid} | Error: {e}\")\n        failed_targets.append(tid)\n\n        # Fallback: linear placeholder — submission file must have no missing rows\n        n      = len(seq)\n        coords = np.zeros((n, 3))\n        for j in range(1, n):\n            coords[j] = coords[j - 1] + [4.0, 0, 0]\n        preds = [coords] * 5\n\n    for j in range(len(seq)):\n        res = {'ID': f\"{tid}_{j+1}\", 'resname': seq[j], 'resid': j + 1}\n        for i in range(5):\n            res[f'x_{i+1}'], res[f'y_{i+1}'], res[f'z_{i+1}'] = preds[i][j]\n        all_predictions.append(res)\n\nprint(f\"\\n✅ Done!  Total time: {time.time() - start_time:.1f}s\")\nprint(f\"❌ Failed sequences: {len(failed_targets)}\")\nif failed_targets:\n    print(f\"   Failed IDs: {failed_targets}\")\n\n# =============================================================================\n# 7. SAVE & VALIDATE\n# =============================================================================\n\nsub  = pd.DataFrame(all_predictions)\ncols = ['ID', 'resname', 'resid'] + [f'{c}_{i}' for i in range(1, 6) for c in ['x', 'y', 'z']]\nsub[cols].to_csv('submission.csv', index=False)\n\nvalidate_submission('submission.csv')\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-22T00:23:43.886158Z","iopub.execute_input":"2026-03-22T00:23:43.886894Z","iopub.status.idle":"2026-03-22T00:26:05.921952Z","shell.execute_reply.started":"2026-03-22T00:23:43.886868Z","shell.execute_reply":"2026-03-22T00:26:05.921336Z"}},"outputs":[],"execution_count":null}]}