{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.11.13","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":31239,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"4c0af182","cell_type":"markdown","source":"\n# 🧬 Stanford RNA 3D Folding Part 2: Advanced Research and Modeling\n\nThis notebook aims to provide a **comprehensive, research-backed, and advanced modeling approach** to the Stanford RNA 3D Folding Part 2 competition. We'll cover:\n\n- An overview of the competition, dataset, and evaluation metric.\n- A summary of recent research in RNA 3D structure prediction and state-of-the-art methods.\n- Exploratory data analysis on the provided training, validation, and test sequences.\n- Implementation of both baseline and advanced models, including secondary-structure-guided folding and ensemble predictions.\n- Example training on a subset of data to illustrate model construction.\n- Generation of a valid submission file for Kaggle.\n\nPlease note that while I implement improved baselines and demonstrate model training pipelines, reproducing the performance of top models like DeepFoldRNA or RhoFold+ requires substantial computational resources and may not be feasible within this notebook. However,In this notebook provides a strong foundation for understanding the problem and building your own advanced models.","metadata":{}},{"id":"7c0990bc","cell_type":"markdown","source":"\n## 🔬 Research Context\n\n### Background: Why RNA Folding is Hard\n\nPredicting the 3D structure of RNA from its sequence is challenging due to the molecule’s flexibility, multiple stable conformations, and limited available data. Unlike proteins, which have more abundant structural data, the Protein Data Bank contains relatively few RNA structures – roughly **1,732 isolated RNA structures** versus **188,726 protein structures**. This scarcity of experimental data makes training deep models difficult.\n\n### Evaluation Metric: TM-score\n\nThe competition uses TM-score (template modeling score) to evaluate predictions. TM-score ranges from 0 to 1, where higher values indicate closer matches to the true structure. A TM-score above 0.5 is often considered a correct fold, while scores below 0.2 indicate poor alignment.\n\n### State-of-the-Art Methods\n\nRecent advances leverage deep learning and language models trained on millions of RNA sequences:\n\n- **RhoFold+**: A deep learning model that uses an RNA language model trained on **~23.7 million sequences** to predict structures. RhoFold+ incorporates secondary structure predictions and interhelical angles, outperforming previous methods and even human experts on benchmark datasets.\n- **DeepFoldRNA**: Combines deep neural networks with multiple sequence alignments and coarse-grained physics-based refinements, achieving high TM-scores on CASP datasets.\n- **Secondary-structure-guided methods**: Many models first predict the RNA secondary structure (e.g., via the Nussinov algorithm) and then generate 3D coordinates by arranging stems and loops according to physics-based constraints.\n\n### Motivation for Baselines\n\nAlthough deep models achieve high accuracy, they require significant computational resources and complex training pipelines. **Good baselines are essential** because they:\n\n- Provide interpretable starting points for model development.\n- Encode physical priors about RNA folding (e.g., A-form helices).\n- Allow quick iteration within the Kaggle environment (CPU or limited GPU)\n- Serve as benchmarks against which more sophisticated methods can be compared.\n","metadata":{}},{"id":"0624948f","cell_type":"markdown","source":"\n## 📂 Dataset Exploration\n\nThe competition provides several CSV files:\n\n- `train_sequences.csv`, `validation_sequences.csv`, and `test_sequences.csv` contain the target RNA sequences and metadata.\n- `train_labels.csv` and `validation_labels.csv` contain experimental C1′ coordinates for each nucleotide in multiple conformations.\n- `sample_submission.csv` gives the expected format of predictions.\n\nI’ll load and inspect these files to understand the sequence lengths and distribution.\n","metadata":{}},{"id":"ed21811b","cell_type":"code","source":"\nimport pandas as pd\nfrom pathlib import Path\n\n# Path to dataset (adjust if running locally on Kaggle)\nDATA_DIR = Path(\"/kaggle/input/stanford-rna-3d-folding-2/\")\n\n# Load sequences\ntrain_sequences = pd.read_csv(DATA_DIR / \"train_sequences.csv\")\nval_sequences = pd.read_csv(DATA_DIR / \"validation_sequences.csv\")\ntest_sequences = pd.read_csv(DATA_DIR / \"test_sequences.csv\")\n\n# Display the number of sequences and a few samples\nprint(\"Train sequences:\", len(train_sequences))\nprint(\"Validation sequences:\", len(val_sequences))\nprint(\"Test sequences:\", len(test_sequences))\n\ntrain_sequences.head()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-08T16:46:16.400737Z","iopub.execute_input":"2026-01-08T16:46:16.401854Z","iopub.status.idle":"2026-01-08T16:46:16.800711Z","shell.execute_reply.started":"2026-01-08T16:46:16.401819Z","shell.execute_reply":"2026-01-08T16:46:16.799578Z"}},"outputs":[],"execution_count":null},{"id":"cf1ea9bd","cell_type":"markdown","source":"\nTo train or evaluate models, we need the labeled 3D coordinates. Each row in `train_labels.csv` corresponds to a single residue with multiple sets of x, y, z coordinates for each experimental conformation.\n","metadata":{}},{"id":"bd6ca604","cell_type":"code","source":"\n# Load labels (only first few rows for inspection)\ntrain_labels = pd.read_csv(DATA_DIR / \"train_labels.csv\", nrows=5000)\nprint(train_labels.head())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-08T16:46:41.229198Z","iopub.execute_input":"2026-01-08T16:46:41.229493Z","iopub.status.idle":"2026-01-08T16:46:41.262212Z","shell.execute_reply.started":"2026-01-08T16:46:41.229472Z","shell.execute_reply":"2026-01-08T16:46:41.261027Z"}},"outputs":[],"execution_count":null},{"id":"a0c3ec10","cell_type":"markdown","source":"\n## 🧱 Baseline Models\n\n### Baseline\n\nThe simplest baseline places the C1′ atoms on an A-form helix with constant rise and twist parameters:\n- **Rise per base**: ~2.8 Å\n- **Twist per base**: ~32.7°\n- **Radius**: ~10 Å\n\nI'll add random noise and random rotations to produce an ensemble of structures. This baseline is fast and easy to implement but fails to capture loops and non-helical elements.\n","metadata":{}},{"id":"264c2f8a","cell_type":"code","source":"\nimport numpy as np\n\nHELIX_RISE = 2.8\nHELIX_TWIST_DEG = 32.7\nHELIX_RADIUS = 10.0\n\ndef helical_fold(sequence: str, seed: int) -> np.ndarray:\n    '''Generate coordinates for a helix with random noise.'''\n    rng = np.random.RandomState(seed)\n    n = len(sequence)\n    coords = np.zeros((n, 3), dtype=np.float32)\n    twist = np.radians(HELIX_TWIST_DEG)\n    for i in range(n):\n        angle = i * twist\n        coords[i] = [\n            HELIX_RADIUS * np.cos(angle),\n            HELIX_RADIUS * np.sin(angle),\n            i * HELIX_RISE\n        ]\n    # Add small random noise\n    coords += rng.normal(0, 0.2, coords.shape)\n    coords -= coords.mean(axis=0)\n    return coords\n\ndef random_rotation(coords: np.ndarray, seed: int) -> np.ndarray:\n    '''Apply a random rotation to 3D coordinates.'''\n    rng = np.random.RandomState(seed)\n    angles = rng.uniform(0, 2*np.pi, size=3)\n    cx, cy, cz = np.cos(angles)\n    sx, sy, sz = np.sin(angles)\n    R = np.array([\n        [cy*cz, -cy*sz, sy],\n        [sx*sy*cz + cx*sz, -sx*sy*sz + cx*cz, -sx*cy],\n        [-cx*sy*cz + sx*sz, cx*sy*sz + sx*cz, cx*cy]\n    ])\n    return (coords - coords.mean(axis=0)) @ R.T\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-08T16:46:52.336398Z","iopub.execute_input":"2026-01-08T16:46:52.337132Z","iopub.status.idle":"2026-01-08T16:46:52.347381Z","shell.execute_reply.started":"2026-01-08T16:46:52.3371Z","shell.execute_reply":"2026-01-08T16:46:52.345682Z"}},"outputs":[],"execution_count":null},{"id":"42001a6e","cell_type":"markdown","source":"\n### Improved with Base-Specific Offsets\n\nTo improve the helix baseline, we incorporate base-specific offsets and small random perturbations. This approach is inspired by the **Boltz1** baseline, which uses canonical A-form geometry and adds base-dependent offsets and noise. The idea is that different nucleotides (A, U, C, G) have slightly different positions relative to an ideal helix. This yields a more realistic ensemble of structures.\n\nBelow is an implementation that generates five conformations for each sequence.\n","metadata":{}},{"id":"e029848a","cell_type":"code","source":"\n# Base-specific offsets for improved helix\nBASE_OFFSETS = {\n    'A': np.array([0.1, 0.1, 0.0]),\n    'U': np.array([-0.1, 0.0, 0.1]),\n    'C': np.array([0.0, -0.1, -0.1]),\n    'G': np.array([0.05, 0.05, -0.05])\n}\n\ndef improved_helical_fold(sequence: str, seed: int) -> np.ndarray:\n    rng = np.random.RandomState(seed)\n    n = len(sequence)\n    coords = np.zeros((n, 3), dtype=np.float32)\n    twist = np.radians(HELIX_TWIST_DEG)\n    for i, base in enumerate(sequence):\n        angle = i * twist\n        # Ideal helix coords\n        coord = np.array([\n            HELIX_RADIUS * np.cos(angle),\n            HELIX_RADIUS * np.sin(angle),\n            i * HELIX_RISE\n        ])\n        # Add base-specific offset\n        offset = BASE_OFFSETS.get(base, np.zeros(3))\n        # Small random noise\n        noise = rng.normal(0, 0.2, 3)\n        coords[i] = coord + offset + noise\n    coords -= coords.mean(axis=0)\n    return coords\n\ndef generate_ensemble(sequence: str, n_models: int = 5, improved: bool = False):\n    models = []\n    for i in range(n_models):\n        if improved:\n            coords = improved_helical_fold(sequence, seed=i * 7919)\n        else:\n            coords = helical_fold(sequence, seed=i * 7919)\n        # Apply random rotation except for the first model\n        if i > 0:\n            coords = random_rotation(coords, seed=i * 13)\n        models.append(coords)\n    return models\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-08T16:47:03.716279Z","iopub.execute_input":"2026-01-08T16:47:03.717194Z","iopub.status.idle":"2026-01-08T16:47:03.726537Z","shell.execute_reply.started":"2026-01-08T16:47:03.717162Z","shell.execute_reply":"2026-01-08T16:47:03.725474Z"}},"outputs":[],"execution_count":null},{"id":"a973a0cf","cell_type":"markdown","source":"\n## 🧠 Secondary-Structure-Guided Modeling\n\n### Nussinov Algorithm for Secondary Structure\n\nMany advanced RNA folding methods first predict the secondary structure (base pairing pattern) before generating 3D coordinates. The **Nussinov algorithm** is a dynamic programming approach that computes a minimum-energy secondary structure by maximizing the number of base pairs. Its computational complexity is \\(O(n^3)\\) with \\(O(n^2)\\) memory. Although Nussinov is simplistic, it serves as a foundation for more advanced thermodynamic models.\n\nBelow I implement a basic Nussinov algorithm and use it to adjust the helical model: paired bases are forced to lie opposite each other in the helix.\n","metadata":{}},{"id":"9338ea69","cell_type":"code","source":"\ndef nussinov_structure(sequence: str) -> set:\n    '''Return base pairs using Nussinov algorithm (maximizes base pairs).'''\n    n = len(sequence)\n    dp = [[0] * n for _ in range(n)]\n    pair = set()\n    # Fill dp table\n    for l in range(1, n):\n        for i in range(n - l):\n            j = i + l\n            # Do not pair\n            best = dp[i+1][j]\n            # Try pairing i-j if complementary\n            if (sequence[i] == 'A' and sequence[j] == 'U') or                (sequence[i] == 'U' and sequence[j] == 'A') or                (sequence[i] == 'C' and sequence[j] == 'G') or                (sequence[i] == 'G' and sequence[j] == 'C'):\n                best = max(best, dp[i+1][j-1] + 1)\n            # Try splitting\n            for k in range(i, j):\n                best = max(best, dp[i][k] + dp[k+1][j])\n            dp[i][j] = best\n    \n    # Traceback to extract pairs\n    def traceback(i, j):\n        if i >= j:\n            return\n        if dp[i][j] == dp[i+1][j]:\n            traceback(i+1, j)\n        else:\n            # Check if i is paired with some k\n            for k in range(i+1, j+1):\n                if (sequence[i] == 'A' and sequence[k] == 'U') or                    (sequence[i] == 'U' and sequence[k] == 'A') or                    (sequence[i] == 'C' and sequence[k] == 'G') or                    (sequence[i] == 'G' and sequence[k] == 'C'):\n                    if dp[i][j] == dp[i+1][k-1] + 1 + dp[k+1][j]:\n                        pair.add((i, k))\n                        traceback(i+1, k-1)\n                        traceback(k+1, j)\n                        return\n            # Try splitting\n            for k in range(i, j):\n                if dp[i][j] == dp[i][k] + dp[k+1][j]:\n                    traceback(i, k)\n                    traceback(k+1, j)\n                    return\n    \n    traceback(0, n-1)\n    return pair\n\n\ndef secondary_structure_guided_fold(sequence: str, seed: int) -> np.ndarray:\n    '''Create coordinates for a sequence where paired bases are placed opposite each other around a helix. Unpaired bases follow a helix with noise.'''\n    pairs = nussinov_structure(sequence)\n    rng = np.random.RandomState(seed)\n    n = len(sequence)\n    coords = np.zeros((n, 3), dtype=np.float32)\n    twist = np.radians(HELIX_TWIST_DEG)\n    used = set()\n    for i in range(n):\n        if i in used:\n            continue\n        # Check if i is paired\n        partner = None\n        for (a, b) in pairs:\n            if a == i:\n                partner = b\n                break\n            if b == i:\n                partner = a\n                break\n        if partner is not None and partner not in used:\n            # Pair i and partner opposite around helix\n            angle = i * twist\n            coord_i = np.array([\n                HELIX_RADIUS * np.cos(angle),\n                HELIX_RADIUS * np.sin(angle),\n                i * HELIX_RISE\n            ])\n            angle_partner = partner * twist + np.pi\n            coord_partner = np.array([\n                HELIX_RADIUS * np.cos(angle_partner),\n                HELIX_RADIUS * np.sin(angle_partner),\n                partner * HELIX_RISE\n            ])\n            coords[i] = coord_i + rng.normal(0, 0.2, 3)\n            coords[partner] = coord_partner + rng.normal(0, 0.2, 3)\n            used.add(i)\n            used.add(partner)\n        else:\n            angle = i * twist\n            coords[i] = [\n                HELIX_RADIUS * np.cos(angle),\n                HELIX_RADIUS * np.sin(angle),\n                i * HELIX_RISE\n            ] + rng.normal(0, 0.2, 3)\n    coords -= coords.mean(axis=0)\n    return coords\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-08T16:47:27.287097Z","iopub.execute_input":"2026-01-08T16:47:27.28796Z","iopub.status.idle":"2026-01-08T16:47:27.304169Z","shell.execute_reply.started":"2026-01-08T16:47:27.287876Z","shell.execute_reply":"2026-01-08T16:47:27.302807Z"}},"outputs":[],"execution_count":null},{"id":"53aba2ee","cell_type":"markdown","source":"\n## 🧑‍🔬 Demonstration: Training a Simple Regression Model\n\nTo illustrate how one might build a machine learning model, I'll train a simple neural network to predict coordinates from sequences. I’ll extract a small subset of sequences and labels due to resource constraints. This example shows the basic workflow: encoding sequences, defining a model, training, and evaluating using TM-score.\n","metadata":{}},{"id":"417e40e7","cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom sklearn.preprocessing import OneHotEncoder\n\n\nsubset_size = 3\nselected_targets = train_sequences.head(subset_size)[\"target_id\"].tolist()\n\nsubset_labels = pd.read_csv(\n    DATA_DIR / \"train_labels.csv\",\n    low_memory=False,\n    dtype={\"chain\": \"string\"}  # optional; safe even if chain exists\n)\n\n\npattern = r\"^(\" + \"|\".join(map(lambda x: str(x).replace(\".\", r\"\\.\"), selected_targets)) + r\")_\"\nsubset_labels = subset_labels[subset_labels[\"ID\"].str.contains(pattern, regex=True)].copy()\n\n# Extract target_id + resid from ID\nsubset_labels[[\"target_id\", \"resid\"]] = subset_labels[\"ID\"].str.extract(r\"^(.+?)_(\\d+)$\")\nsubset_labels[\"resid\"] = subset_labels[\"resid\"].astype(int)\n\n# Build sequence dict for selected targets\nseq_dict = dict(zip(\n    train_sequences[train_sequences[\"target_id\"].isin(selected_targets)][\"target_id\"],\n    train_sequences[train_sequences[\"target_id\"].isin(selected_targets)][\"sequence\"]\n))\n\nencoder = OneHotEncoder(\n    categories=[['A', 'C', 'G', 'U']],\n    sparse_output=False,\n    handle_unknown='ignore'\n)\nencoder.fit(np.array([['A'], ['C'], ['G'], ['U']], dtype=object))\n\nX_list, y_list = [], []\n\nfor target in selected_targets:\n    seq = str(seq_dict[target]).upper().replace(\"T\", \"U\")\n    seq = \"\".join([c for c in seq if c in \"ACGU\"])\n    L = len(seq)\n\n    X_seq = np.array(list(seq), dtype=object).reshape(-1, 1)\n    one_hot = encoder.transform(X_seq).astype(np.float32)\n\n    # Labels: sort by resid to align with sequence positions 1..L\n    lbl = subset_labels[subset_labels[\"target_id\"] == target].sort_values(\"resid\")\n    coords = lbl[[\"x_1\", \"y_1\", \"z_1\"]].to_numpy(dtype=np.float32)\n\n    n = min(len(coords), L)\n    if n == 0:\n        print(f\"Skipping {target}: no label rows found\")\n        continue\n\n    X_list.append(one_hot[:n])\n    y_list.append(coords[:n])\n\nX_np = np.concatenate(X_list, axis=0)\ny_np = np.concatenate(y_list, axis=0)\n\nX_tensor = torch.tensor(X_np, dtype=torch.float32)\ny_tensor = torch.tensor(y_np, dtype=torch.float32)\n\nprint(\"Training data shapes:\", X_tensor.shape, y_tensor.shape)\n\nclass SimpleNN(nn.Module):\n    def __init__(self):\n        super().__init__()\n        self.fc1 = nn.Linear(4, 64)\n        self.fc2 = nn.Linear(64, 32)\n        self.fc3 = nn.Linear(32, 3)\n\n    def forward(self, x):\n        x = torch.relu(self.fc1(x))\n        x = torch.relu(self.fc2(x))\n        return self.fc3(x)\n\nmodel = SimpleNN()\ncriterion = nn.MSELoss()\noptimizer = optim.Adam(model.parameters(), lr=1e-3)\n\nepochs = 50\nfor epoch in range(epochs):\n    optimizer.zero_grad()\n    pred = model(X_tensor)\n    loss = criterion(pred, y_tensor)\n    loss.backward()\n    optimizer.step()\n    if (epoch + 1) % 5 == 0:\n        print(f\"Epoch {epoch+1}/{epochs} | Loss: {loss.item():.6f}\")\n\ndef tm_score(pred_coords: np.ndarray, true_coords: np.ndarray) -> float:\n    pred = pred_coords - pred_coords.mean(axis=0)\n    true = true_coords - true_coords.mean(axis=0)\n    n = len(pred)\n    d0 = 1.24 * ((max(n, 19) - 15) ** (1/3)) - 1.8  # guard small n\n    d0 = max(d0, 0.5)\n    diffs = np.linalg.norm(pred - true, axis=1)\n    return float((1 / n) * np.sum(1 / (1 + (diffs / d0) ** 2)))\n\nwith torch.no_grad():\n    pred_np = model(X_tensor).cpu().numpy()\n\ntm = tm_score(pred_np, y_np)\nprint(f\"Initial TM-score on subset: {tm:.4f}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-08T16:58:52.792848Z","iopub.execute_input":"2026-01-08T16:58:52.793362Z","iopub.status.idle":"2026-01-08T16:59:08.605711Z","shell.execute_reply.started":"2026-01-08T16:58:52.793335Z","shell.execute_reply":"2026-01-08T16:59:08.604715Z"}},"outputs":[],"execution_count":null},{"id":"f2068ec0","cell_type":"markdown","source":"\n## ✅ Conclusion\n\nIn this notebook, we explored the **Stanford RNA 3D Folding Part 2** competition and developed a strong baseline understanding:\n\n- We reviewed recent research showing how deep learning models like **RhoFold+** leverage language models and secondary structure predictions to achieve high accuracy, highlighting the scarcity of RNA structural data and the importance of good baselines.\n- We loaded and explored the provided sequences and labels, analyzing sequence length distributions.\n- We implemented both a **naive helix baseline** and an **improved helix with base-specific offsets**, inspired by the Boltz1 baseline.\n- We incorporated **secondary structure predictions** via the Nussinov algorithm to guide 3D coordinate generation, reflecting how many advanced methods use secondary structure as an intermediate step\n- We demonstrated how to build a **simple neural network** to predict coordinates from sequence one-hot encodings and evaluated it with a simplified TM-score.\n\nThis notebook serves as a robust starting point for your own experimentation. To achieve state-of-the-art performance, consider integrating pre-trained RNA structure models, template search using MSAs, and advanced graph neural networks.\n","metadata":{}},{"id":"c4b3dc80-6288-44c1-b015-ddec9d30f785","cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}