{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","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":false,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Stanford RNA 3D Folding Part 2 – Heuristic Baseline Model\n\nThis notebook implements a **baseline pipeline** for predicting RNA 3D structures from sequence alone, following the submission format required by the *Stanford RNA 3D Folding Part 2* Kaggle competition.\n\nThe goal is to establish a **valid, interpretable baseline**, not a state-of-the-art model.\n\nWe will:\n\n1. Load data and performing EDA  \n2. Define RNA base-pair heuristics  \n3. Generate 3D structures  \n4. Build a Kaggle submission  \n5. Visualize predicted RNA structure  ","metadata":{}},{"cell_type":"markdown","source":"## Imports, Dependencies & Loading dataframe","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom sklearn.decomposition import PCA\nimport matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2026-01-11T08:06:57.760861Z","iopub.execute_input":"2026-01-11T08:06:57.761225Z","iopub.status.idle":"2026-01-11T08:07:01.985966Z","shell.execute_reply.started":"2026-01-11T08:06:57.761198Z","shell.execute_reply":"2026-01-11T08:07:01.984696Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Exploratory Data Analysis and Feature Engineering","metadata":{}},{"cell_type":"code","source":"\n# =========================\n# 0) Constants & utilities\n# =========================\nimport os\nimport numpy as np\nimport pandas as pd\n\nprint(\"packages are loaded\")\n\n# Loading sample data in .CSV format\nsample_dataframe = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/sample_submission.csv\")\n# Loading test sequence data in .CSV format \ntest_dataframe = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\")\n# Loading training sequence data in .CSV format\ntrain_seq_dataframe = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/train_sequences.csv\")\n# Loading validation sequence data in .CSV format\nvalidation_seq_dataframe = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding-2/validation_sequences.csv\")\n\nCANON = {\"A\",\"C\",\"G\",\"U\"}\nBASE_PAIRS = {\n    (\"A\",\"U\"), (\"U\",\"A\"),\n    (\"G\",\"C\"), (\"C\",\"G\"),\n    (\"G\",\"U\"), (\"U\",\"G\"),\n}\n\n# Speed knobs for long sequences (keep runtime under control on Kaggle)\nMAX_STEM_SPAN = 120          # how far j can be from i when searching stems\nMAX_PAIRS_PER_I = 6          # max number of restraints per position\nPAIR_MIN_SEP = 6             # i and j must be at least this far apart\nENSEMBLE_SIZE = 5            # number of models required by competition\n\n# ===================================\n# 1) Fast(er) stem inference (O(n*W))\n# ===================================\ndef infer_stems(sequence, window=4):\n    \"\"\"\n    Mark positions that likely participate in stems using a limited-span search\n    to avoid O(n^2) blow-ups on very long sequences.\n    \"\"\"\n    seq = sequence.upper()\n    n = len(seq)\n    paired = np.zeros(n, dtype=bool)\n    if n < window + 4:\n        return paired\n\n    W = window\n    for i in range(0, n - W):\n        j_start = i + W + 3\n        j_end   = min(n - W, i + W + 3 + MAX_STEM_SPAN)\n        for j in range(j_start, j_end):\n            matches = 0\n            for k in range(W):\n                if (seq[i+k], seq[j-k]) in BASE_PAIRS:\n                    matches += 1\n            if matches >= W - 1:   # NOTE: fixed from &gt;= to >=\n                paired[i:i+W] = True\n                paired[j-W+1:j+1] = True\n    return paired\n\n# ==============================================\n# 2) Sparse \"MSA-like\" restraints (linear-ish)\n# ==============================================\ndef msa_covariance_restraints(seq):\n    \"\"\"\n    Return a sparse list of (i, j) index pairs that are plausible interactions.\n    Caps span and number per i to keep complexity reasonable.\n    \"\"\"\n    seq = seq.upper()\n    n = len(seq)\n    pairs = []\n    for i in range(n):\n        added = 0\n        j_max = min(n, i + MAX_STEM_SPAN)\n        for j in range(i + PAIR_MIN_SEP, j_max):\n            if (seq[i], seq[j]) in BASE_PAIRS:\n                pairs.append((i, j))\n                added += 1\n                if added >= MAX_PAIRS_PER_I:\n                    break\n    return pairs\n\n# ======================================\n# 3) Lightweight geometric baseline core\n# ======================================\ndef gc_radius(base):\n    return 5.5 if base in {\"G\",\"C\"} else 7.5\n\ndef template_fragment(seq, start_idx):\n    # tiny fragment to perturb backbone; bounded by seq end\n    frag_length = min(5, len(seq) - start_idx)\n    if frag_length <= 0:\n        return np.zeros((0,3))\n    t = np.linspace(0, np.pi, frag_length)\n    return np.stack([np.cos(t), np.sin(t), t], axis=1)\n\ndef generate_structure(seq, seed=0):\n    \"\"\"\n    Produce an Nx3 coordinate array for RNA C1' positions.\n    Designed to be fast and always return (len(seq), 3).\n    \"\"\"\n    if not isinstance(seq, str):\n        seq = str(seq)\n    seq = seq.strip().upper()\n    n = len(seq)\n\n    # Guard: empty/invalid sequences\n    if n == 0 or not set(seq).issubset(CANON):\n        return np.zeros((max(n,1), 3), dtype=float)\n\n    rng = np.random.default_rng(seed)\n\n    # 3a) helix-like backbone with sequence-dependent radius & optional stem squeeze\n    t = np.linspace(0, 2*np.pi*n/10, n)\n    paired = infer_stems(seq)  # boolean mask\n    coords = np.zeros((n,3), dtype=float)\n    for i, base in enumerate(seq):\n        r = gc_radius(base)\n        if paired[i]:\n            r *= 0.85\n        coords[i] = [r*np.cos(t[i]), r*np.sin(t[i]), 3.0*t[i]]\n\n    # 3b) add tiny template fragments every 10 residues (bounded)\n    for i in range(0, n, 10):\n        frag = template_fragment(seq, i)\n        if frag.shape[0] > 0:\n            coords[i:i+frag.shape[0]] += frag\n\n    # 3c) apply sparse restraints (pull together plausible pairs)\n    for (i, j) in msa_covariance_restraints(seq):\n        diff = coords[j] - coords[i]\n        coords[j] -= 0.05 * diff\n\n    # 3d) small Gaussian noise for diversity\n    coords += rng.normal(0, 0.25, size=coords.shape)\n\n    # 3e) center structure (zero-mean)\n    coords -= coords.mean(axis=0, keepdims=True)\n    return coords\n\n# ========================================\n# 4) Build the competition submission file\n# ========================================\ndef build_submission(test_df, out_csv=\"submission.csv\"):\n    \"\"\"\n    Create submission.csv with columns:\n    ID, resname, resid, x_1,y_1,z_1, ..., x_5,y_5,z_5\n    \"\"\"\n    rows = []\n    seeds = list(range(ENSEMBLE_SIZE))\n\n    # fixed column order required by Kaggle\n    cols = [\"ID\", \"resname\", \"resid\"] + \\\n           [f\"{axis}_{m}\" for m in range(1, ENSEMBLE_SIZE+1) for axis in (\"x\",\"y\",\"z\")]\n\n    for _, r in test_df.iterrows():\n        tid = str(r[\"target_id\"])\n        seq = str(r[\"sequence\"]).strip().upper()\n\n        # Validate sequence is canonical; if not, map T->U and filter others\n        if not set(seq).issubset(CANON):\n            seq = seq.replace(\"T\",\"U\")\n            seq = \"\".join(ch for ch in seq if ch in \"ACGU\")\n\n        # Generate 5 models once per target\n        models = [generate_structure(seq, seed=s) for s in seeds]\n\n        # Safety: ensure each model has correct shape\n        n = len(seq)\n        ok = all(m.shape == (n,3) for m in models)\n        if not ok:\n            models = [np.resize(m, (n,3)) for m in models]\n\n        # emit per-residue rows\n        for i, nt in enumerate(seq, start=1):\n            row = {\"ID\": f\"{tid}_{i}\", \"resname\": nt, \"resid\": i}\n            for m_idx, M in enumerate(models, start=1):\n                row[f\"x_{m_idx}\"] = float(M[i-1, 0])\n                row[f\"y_{m_idx}\"] = float(M[i-1, 1])\n                row[f\"z_{m_idx}\"] = float(M[i-1, 2])\n            rows.append(row)\n\n    sub = pd.DataFrame(rows, columns=cols)\n    sub.to_csv(out_csv, index=False)\n    return sub\n\n# =========================================\n# 5) Run once and verify the file is saved\n# =========================================\nsubmission = build_submission(test_dataframe, out_csv=\"submission.csv\")\n\n# show head and confirm presence in /kaggle/working\nprint(submission.head(3))\nprint(\"\\nFiles in /kaggle/working:\", os.listdir(\"/kaggle/working\"))\nassert os.path.exists(\"/kaggle/working/submission.csv\"), \"submission.csv was not written!\"\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T08:07:01.988104Z","iopub.execute_input":"2026-01-11T08:07:01.988579Z","iopub.status.idle":"2026-01-11T08:07:09.232856Z","shell.execute_reply.started":"2026-01-11T08:07:01.98855Z","shell.execute_reply":"2026-01-11T08:07:09.231728Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Visualization of predicted 3D model of RNA","metadata":{}},{"cell_type":"code","source":"import matplotlib.pyplot as plt\nfrom mpl_toolkits.mplot3d import Axes3D\nimport seaborn as sns\nfrom sklearn.decomposition import PCA\n\ndef plot_rna_structure(coords, seq, title=\"RNA 3D Structure\"):\n    if coords is None or len(coords) != len(seq):\n        print(\"Error: Invalid coordinates. Cannot plot structure.\")\n        return\n    \n    fig = plt.figure(figsize=(8,6))\n    ax = fig.add_subplot(111, projection='3d')\n\n    x, y, z = coords[:,0], coords[:,1], coords[:,2]\n\n    ax.plot(x, y, z, linewidth=1.5, color='blue')\n    ax.scatter(x, y, z, s=25, color='red')\n\n    for i, base in enumerate(seq):\n        ax.text(x[i], y[i], z[i], base, fontsize=8)\n\n    ax.set_title(title)\n    ax.set_xlabel(\"X (Å)\")\n    ax.set_ylabel(\"Y (Å)\")\n    ax.set_zlabel(\"Z (Å)\")\n    plt.show()\n\n# Example usage\nexample_seq = test_dataframe.iloc[0][\"sequence\"]\ncoords_example = generate_structure(example_seq, seed=0)\nplot_rna_structure(coords_example, example_seq, \"Predicted RNA Structure (Seed 0)\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-11T08:07:09.234213Z","iopub.execute_input":"2026-01-11T08:07:09.234566Z","iopub.status.idle":"2026-01-11T08:07:09.990641Z","shell.execute_reply.started":"2026-01-11T08:07:09.234531Z","shell.execute_reply":"2026-01-11T08:07:09.989591Z"}},"outputs":[],"execution_count":null}]}