{"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":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14604295,"datasetId":9328538,"databundleVersionId":15440074}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":155.274731,"end_time":"2026-02-03T09:12:36.032588","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-02-03T09:10:00.757857","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"643f6782","cell_type":"markdown","source":"# 🧬 RNA 3D Structure Prediction\n**Competition:** Stanford RNA 3D Folding 2  \n**Approach:** Homology-based template matching with adaptive geometry refinement  \n","metadata":{}},{"id":"a586fdfb","cell_type":"markdown","source":"## 1. Environment Setup\n","metadata":{}},{"id":"56bfd799","cell_type":"code","source":"import subprocess, sys, os, numpy as np\n\nWHEEL = (\n    '/kaggle/input/datasets/kami1976/biopython-cp312/'\n    'biopython-1.86-cp312-cp312-manylinux2014_x86_64'\n    '.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl'\n)\n\nif os.path.exists(WHEEL):\n    subprocess.check_call([sys.executable, '-m', 'pip', 'install', '--no-index', WHEEL])\n    print('BioPython installed from offline wheel.')\nelse:\n    raise FileNotFoundError(\n        f\"Wheel not found at {WHEEL}\\n\"\n        \"Please attach the dataset: kami1976/biopython-cp312\"\n    )\n\nfrom Bio.Align import PairwiseAligner\n\naligner = PairwiseAligner()\naligner.mode             = 'global'\naligner.match_score      =  2\naligner.mismatch_score   = -1.5\naligner.open_gap_score   = -8\naligner.extend_gap_score = -0.4\nfor side in ('query_left','query_right','target_left','target_right'):\n    setattr(aligner, f'{side}_open_gap_score',   -8)\n    setattr(aligner, f'{side}_extend_gap_score', -0.4)\n\nprint(\"BioPython PairwiseAligner ready.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:19:39.602028Z","iopub.execute_input":"2026-02-21T06:19:39.602754Z","iopub.status.idle":"2026-02-21T06:19:44.671347Z","shell.execute_reply.started":"2026-02-21T06:19:39.602724Z","shell.execute_reply":"2026-02-21T06:19:44.670594Z"}},"outputs":[],"execution_count":null},{"id":"164897b7","cell_type":"markdown","source":"## 2. Imports & Data Loading\nLoad train/test sequences and ground-truth coordinates.  \nWe also attach a robust FASTA parser and stoichiometry utilities.","metadata":{}},{"id":"0ae06f07","cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport random\nimport time\nimport warnings\nimport os, sys\n\nwarnings.filterwarnings('ignore')\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\"))\nprint(f\"Train sequences : {len(train_seqs)}\")\nprint(f\"Test  sequences : {len(test_seqs)}\")\nprint(f\"Train labels    : {len(train_labels)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:20:19.113459Z","iopub.execute_input":"2026-02-21T06:20:19.114195Z","iopub.status.idle":"2026-02-21T06:20:28.630184Z","shell.execute_reply.started":"2026-02-21T06:20:19.114165Z","shell.execute_reply":"2026-02-21T06:20:28.629559Z"}},"outputs":[],"execution_count":null},{"id":"85ec2692","cell_type":"markdown","source":"## 3. FASTA Parser & Stoichiometry Utilities\nParse multi-chain sequences and stoichiometry strings (e.g. `A:2;B:1`) to\nidentify per-chain residue segments in the full concatenated sequence.","metadata":{}},{"id":"d3f6f76f","cell_type":"code","source":"# --- Robust import for Kaggle's extra/parse_fasta_py.py ---\ntry:\n    import typing as _typing\n    import builtins as _builtins\n    _builtins.Dict  = getattr(_typing, \"Dict\")\n    _builtins.Tuple = getattr(_typing, \"Tuple\")\n    _builtins.List  = getattr(_typing, \"List\")\n    from parse_fasta_py import parse_fasta as _parse_fasta_raw\n    def parse_fasta(fasta_content: str):\n        d = _parse_fasta_raw(fasta_content)\n        out = {}\n        for k, v in d.items():\n            out[k] = v[0] if isinstance(v, tuple) else v\n        return out\nexcept Exception:\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: continue\n            if line.startswith(\">\"):\n                if cur is not None: out[cur] = \"\".join(seq_parts)\n                header = line[1:]\n                cur = header.split()[0]\n                seq_parts = []\n            else: seq_parts.append(line.replace(\" \", \"\"))\n        if cur is not None: out[cur] = \"\".join(seq_parts)\n        return out\n\ndef parse_stoichiometry(stoich: str):\n    if pd.isna(stoich) or str(stoich).strip() == \"\": 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\ndef get_chain_segments(row):\n    seq    = row['sequence']\n    stoich = row.get('stoichiometry', '')\n    all_seq = row.get('all_sequences', '')\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    try:\n        chain_dict = parse_fasta(all_seq)\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: return [(0, len(seq))]\n            for _ in range(cnt):\n                segs.append((pos, pos + len(base)))\n                pos += len(base)\n        if pos != len(seq): return [(0, len(seq))]\n        return segs\n    except Exception: return [(0, len(seq))]\n\ndef build_segments_map(df):\n    seg_map, 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\ntrain_segs_map, train_stoich_map = build_segments_map(train_seqs)\ntest_segs_map,  test_stoich_map  = build_segments_map(test_seqs)\n\ndef process_labels(labels_df):\n    coords_dict = {}\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\ntrain_coords_dict = process_labels(train_labels)\nprint(f\"Loaded {len(train_coords_dict)} training coordinate sets\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:20:43.187742Z","iopub.execute_input":"2026-02-21T06:20:43.188405Z","iopub.status.idle":"2026-02-21T06:20:55.710577Z","shell.execute_reply.started":"2026-02-21T06:20:43.188378Z","shell.execute_reply":"2026-02-21T06:20:55.709709Z"}},"outputs":[],"execution_count":null},{"id":"0af0f847","cell_type":"markdown","source":"## 4. Sequence Alignment (BioPython PairwiseAligner)\nWe use **global alignment** with strong gap penalties to prevent residue index\nsliding — a critical requirement because the evaluation uses TM-score which is\nsensitive to residue correspondence.","metadata":{}},{"id":"f50b97cb","cell_type":"code","source":"# Aligner configured in §1. Quick sanity check:\nprint(f\"match={aligner.match_score}  mismatch={aligner.mismatch_score}\")\nprint(f\"gap open={aligner.open_gap_score}  gap extend={aligner.extend_gap_score}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:21:02.334357Z","iopub.execute_input":"2026-02-21T06:21:02.335207Z","iopub.status.idle":"2026-02-21T06:21:02.339235Z","shell.execute_reply.started":"2026-02-21T06:21:02.33518Z","shell.execute_reply":"2026-02-21T06:21:02.338528Z"}},"outputs":[],"execution_count":null},{"id":"990e5782","cell_type":"markdown","source":"## 5. Template Search & Coordinate Transfer\nFor each test sequence:\n1. **`find_similar_sequences`** — score all training templates via length-filtered\n   fast alignment scoring and return the top-N candidates.\n2. **`adapt_template_to_query`** — vectorised alignment mapping transfers 3-D\n   coordinates from template to query positions; gaps are filled by linear\n   interpolation.","metadata":{}},{"id":"eb0ac1fb","cell_type":"code","source":"def find_similar_sequences(query_seq, train_seqs_df, train_coords_dict, top_n=5):\n    similar_seqs = []\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: continue\n        if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > 0.3: continue\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    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords):\n    alignment = next(iter(aligner.align(query_seq, template_seq)))\n    new_coords = np.full((len(query_seq), 3), np.nan)\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):\n            new_coords[q_start:q_end] = t_chunk\n    # Interpolate / extrapolate unmapped residues\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            if prev_v >= 0 and next_v >= 0:\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: new_coords[i] = new_coords[prev_v] + [3, 0, 0]\n            elif next_v >= 0: new_coords[i] = new_coords[next_v] + [3, 0, 0]\n            else:             new_coords[i] = [i*3, 0, 0]\n    return np.nan_to_num(new_coords)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:21:07.552939Z","iopub.execute_input":"2026-02-21T06:21:07.553582Z","iopub.status.idle":"2026-02-21T06:21:07.563385Z","shell.execute_reply.started":"2026-02-21T06:21:07.553553Z","shell.execute_reply":"2026-02-21T06:21:07.562743Z"}},"outputs":[],"execution_count":null},{"id":"73113f9c","cell_type":"markdown","source":"## 6. Adaptive Geometry Refinement\nAfter coordinate transfer we apply within-chain physical constraints:\n- **Bond correction** — nudge consecutive residue pairs toward ~5.95 Å\n- **i,i+2 distance** — soft target ~10.2 Å to enforce local backbone curvature\n- **Laplacian smoothing** — removes sharp kinks\n- **Self-avoidance** — repels residues closer than 3.2 Å (for long chains)\n\nCorrection strength scales inversely with template similarity confidence.","metadata":{}},{"id":"2c715976","cell_type":"code","source":"def adaptive_rna_constraints(coordinates, target_id, confidence=1.0, passes=2):\n    coords = coordinates.copy()\n    segments = test_segs_map.get(target_id, [(0, len(coords))])\n    strength = max(0.75 * (1.0 - min(confidence, 0.97)), 0.02)\n    for _ in range(passes):\n        for (s, e) in segments:\n            X = coords[s:e]\n            L = e - s\n            if L < 3: coords[s:e] = X; continue\n            # (1) bond i,i+1 → ~5.95Å\n            d = X[1:] - X[:-1]\n            dist = np.linalg.norm(d, axis=1) + 1e-6\n            adj = (d * ((5.95 - dist)/dist)[:, None]) * (0.22 * strength)\n            X[:-1] -= adj;  X[1:] += adj\n            # (2) i,i+2 → ~10.2Å\n            d2 = X[2:] - X[:-2]\n            dist2 = np.linalg.norm(d2, axis=1) + 1e-6\n            adj2 = (d2 * ((10.2 - dist2)/dist2)[:, None]) * (0.10 * strength)\n            X[:-2] -= adj2; X[2:] += adj2\n            # (3) Laplacian smoothing\n            lap = 0.5*(X[:-2] + X[2:]) - X[1:-1]\n            X[1:-1] += (0.06 * strength) * lap\n            # (4) Self-avoidance (large chains only)\n            if L >= 25:\n                k = min(L, 160) if L > 220 else L\n                idx = np.linspace(0, L-1, k).astype(int) if k < L else np.arange(L)\n                P = X[idx]\n                diff = P[:, None, :] - P[None, :, :]\n                distm = np.linalg.norm(diff, axis=2) + 1e-6\n                sep   = np.abs(idx[:, None] - idx[None, :])\n                mask  = (sep > 2) & (distm < 3.2)\n                if np.any(mask):\n                    force = (3.2 - distm) / distm\n                    vec   = (diff * force[:,:,None] * mask[:,:,None]).sum(axis=1)\n                    X[idx] += (0.015 * strength) * vec\n            coords[s:e] = X\n    return coords","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:21:13.671673Z","iopub.execute_input":"2026-02-21T06:21:13.671936Z","iopub.status.idle":"2026-02-21T06:21:13.681872Z","shell.execute_reply.started":"2026-02-21T06:21:13.671917Z","shell.execute_reply":"2026-02-21T06:21:13.681161Z"}},"outputs":[],"execution_count":null},{"id":"cc303ac3","cell_type":"markdown","source":"## 7. Structural Diversity Generators\nThe competition evaluates 5 predictions per target (best-of-5 TM-score).  \nWe generate diversity through three complementary perturbations:\n- **Hinge rotation** — rigid rotation of one chain segment around a pivot\n- **Chain jitter** — independent rigid-body rotation + translation per chain\n- **Smooth wiggle** — low-frequency spline-interpolated displacement field","metadata":{}},{"id":"dd5036bc","cell_type":"code","source":"def _rotmat(axis, ang):\n    axis = np.asarray(axis, float)\n    axis = axis / (np.linalg.norm(axis) + 1e-12)\n    x, y, z = axis\n    c, s = np.cos(ang), np.sin(ang); C = 1.0 - c\n    return np.array([[c+x*x*C, x*y*C-z*s, x*z*C+y*s],\n                      [y*x*C+z*s, c+y*y*C, y*z*C-x*s],\n                      [z*x*C-y*s, z*y*C+x*s, c+z*z*C]], dtype=float)\n\ndef apply_hinge(coords, seg, rng, max_angle_deg=25):\n    s, e = seg; L = e - s\n    if L < 30: return coords\n    pivot = s + int(rng.integers(10, L - 10))\n    R = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-max_angle_deg, max_angle_deg))))\n    X = coords.copy(); p0 = X[pivot].copy()\n    X[pivot+1:e] = (X[pivot+1:e] - p0) @ R.T + p0\n    return X\n\ndef jitter_chains(coords, segments, rng, max_angle_deg=12, max_trans=1.5):\n    X = coords.copy()\n    global_center = X.mean(axis=0, keepdims=True)\n    for (s, e) in segments:\n        R = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-max_angle_deg, max_angle_deg))))\n        shift = rng.normal(size=3); shift = shift/(np.linalg.norm(shift)+1e-12)*float(rng.uniform(0,max_trans))\n        c = X[s:e].mean(axis=0, keepdims=True)\n        X[s:e] = (X[s:e] - c) @ R.T + c + shift\n    X -= X.mean(axis=0, keepdims=True) - global_center\n    return X\n\ndef smooth_wiggle(coords, segments, rng, amp=0.8):\n    X = coords.copy()\n    for (s, e) in segments:\n        L = e - s\n        if L < 20: continue\n        ctrl_x = np.linspace(0, L-1, 6)\n        ctrl_disp = rng.normal(0, amp, size=(6, 3))\n        t = np.arange(L)\n        disp = np.vstack([np.interp(t, ctrl_x, ctrl_disp[:, k]) for k in range(3)]).T\n        X[s:e] += disp\n    return X","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:21:18.263205Z","iopub.execute_input":"2026-02-21T06:21:18.263523Z","iopub.status.idle":"2026-02-21T06:21:18.274249Z","shell.execute_reply.started":"2026-02-21T06:21:18.263498Z","shell.execute_reply":"2026-02-21T06:21:18.273586Z"}},"outputs":[],"execution_count":null},{"id":"9e021841","cell_type":"markdown","source":"## 8. Main Prediction Pipeline\n`predict_rna_structures` orchestrates the full workflow per target:\n1. Retrieve top-30 template candidates via fast alignment scoring\n2. Generate **5 diverse predictions** using the transforms above\n3. Refine each prediction with the adaptive geometry constraints","metadata":{}},{"id":"8e60c285","cell_type":"code","source":"def predict_rna_structures(row, train_seqs_df, train_coords_dict, n_predictions=5):\n    tid = row['target_id']\n    seq = row['sequence']\n    assert set(seq).issubset(set(\"ACGU\")), f\"Non-ACGU in {tid}\"\n    segments = test_segs_map.get(tid, [(0, len(seq))])\n    cands = find_similar_sequences(seq, train_seqs_df, train_coords_dict, top_n=30)\n    assert all(len(c[3]) == len(c[1]) for c in cands)\n    predictions = []; used = set()\n    for i in range(n_predictions):\n        seed = (abs(hash(tid)) + i * 10007) % (2**32)\n        rng = np.random.default_rng(seed)\n        if not cands:\n            coords = np.zeros((len(seq), 3), dtype=float)\n            for (s, e) in segments:\n                for j in range(s+1, e): coords[j] = coords[j-1] + [5.95, 0, 0]\n            predictions.append(coords); continue\n        if i == 0:\n            t_id, t_seq, sim, t_coords = cands[0]\n        else:\n            K = min(12, len(cands))\n            sims = np.array([cands[k][2] for k in range(K)], float)\n            w = np.exp((sims - sims.max()) / 0.08)\n            for k in range(K):\n                if cands[k][0] in used: w[k] *= 0.10\n            w = w / (w.sum() + 1e-12)\n            k = int(rng.choice(np.arange(K), p=w))\n            t_id, t_seq, sim, t_coords = cands[k]\n        used.add(t_id)\n        adapted = adapt_template_to_query(seq, t_seq, t_coords)\n        if   i == 0: X = adapted\n        elif i == 1: X = adapted + rng.normal(0, max(0.01, (0.40-sim)*0.06), adapted.shape)\n        elif i == 2: X = apply_hinge(adapted, max(segments, key=lambda se: se[1]-se[0]), rng, 22)\n        elif i == 3: X = jitter_chains(adapted, segments, rng, 10, 1.0)\n        else:        X = smooth_wiggle(adapted, segments, rng, 0.7)\n        predictions.append(adaptive_rna_constraints(X, tid, confidence=sim, passes=2))\n    return predictions","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:21:22.995141Z","iopub.execute_input":"2026-02-21T06:21:22.995711Z","iopub.status.idle":"2026-02-21T06:21:23.005443Z","shell.execute_reply.started":"2026-02-21T06:21:22.995682Z","shell.execute_reply":"2026-02-21T06:21:23.004734Z"}},"outputs":[],"execution_count":null},{"id":"569643e3","cell_type":"markdown","source":"## 9. Generate Submission\nIterate over all test targets, produce 5 coordinate sets each, then write\n`submission.csv` in the competition-required format.","metadata":{}},{"id":"30bb8454","cell_type":"code","source":"all_predictions = []\nstart_time = time.time()\nfor idx, row in test_seqs.iterrows():\n    if idx % 10 == 0: print(f\"Processing {idx} | {time.time()-start_time:.1f}s\")\n    tid, seq = row['target_id'], row['sequence']\n    preds = predict_rna_structures(row, train_seqs, train_coords_dict)\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\nsub = pd.DataFrame(all_predictions)\ncols = ['ID', 'resname', 'resid'] + [f'{c}_{i}' for i in range(1,6) for c in ['x','y','z']]\ncoord_cols = [c for c in cols if c.startswith(('x_','y_','z_'))]\nsub[coord_cols] = sub[coord_cols].clip(-999.999, 9999.999)\nsub[cols].to_csv('submission.csv', index=False)\nprint(\"submission.csv saved!\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-21T06:21:27.142119Z","iopub.execute_input":"2026-02-21T06:21:27.142738Z","iopub.status.idle":"2026-02-21T06:23:32.520534Z","shell.execute_reply.started":"2026-02-21T06:21:27.142709Z","shell.execute_reply":"2026-02-21T06:23:32.519842Z"}},"outputs":[],"execution_count":null},{"id":"e9dfd666-6b3c-496a-9101-2a9680043865","cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}