{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.x"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":6233380,"sourceType":"datasetVersion","datasetId":3580819},{"sourceId":12248008,"sourceType":"datasetVersion","datasetId":7717350}],"dockerImageVersionId":31236,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"1708673b","cell_type":"markdown","source":"# Stanford RNA 3D Folding Part 2 (Kaggle)  \n## A reproducible, research style notebook for high accuracy RNA 3D coordinate prediction\n\n**Author:** Olaf Yunus Laitinen Imanov  \n**Target environment:** Kaggle Code Competition (internet disabled at run time)  \n**Recommended accelerator:** **GPU (NVIDIA Tesla P100, 16 GB)**  \n**Objective:** Given an RNA sequence, predict the 3D coordinates of the **C1' atom** for each nucleotide, and output **five** candidate structures per target.\n\n### Abstract\nRibonucleic acid (RNA) structure is strongly coupled to function, yet RNA 3D structure prediction remains challenging due to sparse experimental data, conformational heterogeneity, and complex long range interactions. This notebook presents a complete end to end pipeline for the Stanford RNA 3D Folding Part 2 competition. The pipeline is built around **RhoFold+**, a modern deep learning system for RNA 3D structure prediction that leverages language model representations and optional evolutionary information from multiple sequence alignments. Inference is constrained to the Kaggle execution model: no internet access, time limits, and a requirement to emit a `submission.csv` that contains five coordinate sets per test target.\n\nThe notebook is structured like a scientific manuscript: problem definition, data specification, metric derivation, computational considerations, methods, and an implementation that produces a valid submission. To support strict reproducibility, all randomness is controlled, intermediate artifacts are cached, and runtime critical sections are profiled. Where model weights are required, the notebook uses only datasets that can be attached through Kaggle's \"Add data\" mechanism.\n\n### Key contributions\n1. A rigorous specification of the data layout, identifiers, and edge cases for multi chain targets.\n2. A reference implementation of the competition metric (TM score) and its alignment assumptions.\n3. A fast inference workflow for RhoFold+ with robust extraction of **C1'** coordinates and format validation.\n4. A fallback template style baseline to guarantee a valid submission if external weights are unavailable.\n\n### References (selected, see full list at the end)\n- Kaggle competition page for Stanford RNA 3D Folding Part 2.  \n- US align and TM score objective function for structure alignment.  \n- RhoFold+ repository and documentation for inference inputs and outputs.  \n- Stanford RNA 3D Folding Part 1 outcome preprint highlighting template based pipelines.\n\n","metadata":{}},{"id":"7f68176b","cell_type":"code","source":"# Environment and reproducibility setup\nimport os, sys, json, math, time, random, re, glob, subprocess\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\n\nSEED = 20260109\nrandom.seed(SEED)\nnp.random.seed(SEED)\n\npd.set_option(\"display.max_columns\", 200)\npd.set_option(\"display.width\", 200)\n\nWORKDIR = Path(\"/kaggle/working\") if Path(\"/kaggle/working\").exists() else Path(\".\")\nINPUT_ROOT = Path(\"/kaggle/input\") if Path(\"/kaggle/input\").exists() else Path(\".\")\n\nprint(\"WORKDIR:\", WORKDIR)\nprint(\"INPUT_ROOT:\", INPUT_ROOT)\n\ndef run(cmd):\n    print(\"$\", \" \".join(cmd))\n    p = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, text=True)\n    print(p.stdout[:4000])\n    return p.returncode, p.stdout\n\n# CPU and GPU introspection (safe even if no GPU is enabled)\nrun([\"bash\", \"-lc\", \"uname -a\"])\nrun([\"bash\", \"-lc\", \"python -V\"])\nrun([\"bash\", \"-lc\", \"nvidia-smi || true\"])\nrun([\"bash\", \"-lc\", \"lscpu | head -n 40 || true\"])","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:46:17.944108Z","iopub.execute_input":"2026-01-09T18:46:17.944703Z","iopub.status.idle":"2026-01-09T18:46:18.030575Z","shell.execute_reply.started":"2026-01-09T18:46:17.944666Z","shell.execute_reply":"2026-01-09T18:46:18.029841Z"}},"outputs":[],"execution_count":null},{"id":"e897d006","cell_type":"markdown","source":"# 1. Competition formulation and notation\n\nLet the RNA sequence for a target be\n$$\n\\mathbf{s} = (s_1, s_2, \\dots, s_L), \\qquad s_i \\in \\{A,C,G,U\\},\n$$\nwith length $L$. The required prediction is a set of Cartesian coordinates for the **C1' atom** of each residue:\n$$\n\\mathbf{X} = (\\mathbf{x}_1, \\mathbf{x}_2, \\dots, \\mathbf{x}_L), \\qquad \\mathbf{x}_i \\in \\mathbb{R}^3.\n$$\nThe submission format requires **five** candidate structures\n$\\{\\mathbf{X}^{(k)}\\}_{k=1}^5$ for every target sequence.\n\nThis competition evaluates predictions using **TM score** after an optimal rigid alignment (rotation plus translation) computed by the **US align** tool. TM score is bounded in $[0,1]$ with higher values indicating better global topology agreement. The Part 2 rules introduce an additional constraint: the alignment rewards only residues that match the reference residue numbering. Practically, this means that if you output coordinates for residue $i$, the scoring aligns that position to residue $i$ in the reference and does not allow flexible renumbering.\n\nWe will therefore treat the learning and inference problem as a sequence to structure mapping\n$$\nf_\\theta : \\{A,C,G,U\\}^L \\to (\\mathbb{R}^3)^L,\n$$\nand pay particular attention to:\n- strict residue indexing (1 based)\n- target identifier formatting in `ID = target_id + \"_\" + resid`\n- coordinate clipping required by legacy PDB formatting\n\nA final practical constraint is that Kaggle notebooks for this competition run with internet disabled. All model weights and auxiliary assets must be shipped either with the competition dataset itself or as attached public Kaggle datasets.\n","metadata":{}},{"id":"b5345da8","cell_type":"code","source":"# Competition dataset root paths (as provided by the user)\nCOMP = INPUT_ROOT / \"stanford-rna-3d-folding-2\"\n\npaths = {\n    \"sample_submission\": COMP / \"sample_submission.csv\",\n    \"train_sequences\": COMP / \"train_sequences.csv\",\n    \"train_labels\": COMP / \"train_labels.csv\",\n    \"validation_sequences\": COMP / \"validation_sequences.csv\",\n    \"validation_labels\": COMP / \"validation_labels.csv\",\n    \"test_sequences\": COMP / \"test_sequences.csv\",\n    \"MSA_dir\": COMP / \"MSA\",\n    \"PDB_RNA_dir\": COMP / \"PDB_RNA\",\n    \"extra_dir\": COMP / \"extra\",\n    \"extra_readme\": COMP / \"extra\" / \"README.md\",\n    \"parse_fasta_py\": COMP / \"extra\" / \"parse_fasta_py.py\",\n    \"rna_metadata\": COMP / \"extra\" / \"rna_metadata.csv\",\n}\nfor k, v in paths.items():\n    print(f\"{k:20s} -> {v}  exists={v.exists()}\")\n\n# List top level competition input directory\nif COMP.exists():\n    print(\"\\nTop level files and folders:\")\n    for p in sorted(COMP.iterdir()):\n        print(\" -\", p.name)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:46:18.032086Z","iopub.execute_input":"2026-01-09T18:46:18.032298Z","iopub.status.idle":"2026-01-09T18:46:18.048451Z","shell.execute_reply.started":"2026-01-09T18:46:18.032279Z","shell.execute_reply":"2026-01-09T18:46:18.047756Z"}},"outputs":[],"execution_count":null},{"id":"3a177502","cell_type":"markdown","source":"# 2. Data schema and identifier mechanics\n\nThe competition provides separate CSV files for sequences and labels. The essential fields are summarized here and then verified directly from the provided files in the accompanying code cell.\n\n## 2.1 File inventory (competition dataset)\n\n| Path (read only) | Artifact | Purpose |\n|---|---|---|\n| `/kaggle/input/stanford-rna-3d-folding-2/train_sequences.csv` | Train sequences | Target ids and RNA sequences for training |\n| `/kaggle/input/stanford-rna-3d-folding-2/train_labels.csv` | Train labels | Experimental C1' coordinates for training targets |\n| `/kaggle/input/stanford-rna-3d-folding-2/validation_sequences.csv` | Validation sequences | Public validation targets |\n| `/kaggle/input/stanford-rna-3d-folding-2/validation_labels.csv` | Validation labels | Experimental coordinates for validation targets |\n| `/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv` | Test sequences | Public test targets used for submissions |\n| `/kaggle/input/stanford-rna-3d-folding-2/sample_submission.csv` | Submission template | Exact output schema and row order |\n| `/kaggle/input/stanford-rna-3d-folding-2/MSA/` | MSAs | Multiple sequence alignments for targets |\n| `/kaggle/input/stanford-rna-3d-folding-2/PDB_RNA/` | Structures | mmCIF files for RNA containing PDB entries |\n| `/kaggle/input/stanford-rna-3d-folding-2/extra/` | Extras | Metadata, README, helper scripts |\n\n## 2.2 Sequence tables\nFor each split (`train_sequences.csv`, `validation_sequences.csv`, `test_sequences.csv`), the table includes:\n- `target_id`: identifier for the RNA target\n- `sequence`: concatenated RNA sequence for the chains specified by `stoichiometry`\n- `stoichiometry`: a semicolon separated mapping like `{chain:number}`\n- `all_sequences`: FASTA formatted sequences for all molecular chains present in the experimentally solved structure, including partners\n- `temporal_cutoff`: date associated with release or publication\n- optional ligand fields: `ligand_ids`, `ligand_SMILES`\n\n## 2.3 Label tables\nFor labeled splits (`train_labels.csv`, `validation_labels.csv`) the label file contains:\n- `ID`: `target_id_resid` where resid is 1 based\n- `resname`: nucleotide letter at that position\n- `resid`: residue index (1 based)\n- `(x_j, y_j, z_j)`: C1' coordinates for one or more experimentally observed structures\n\nSome targets have multiple experimentally captured conformations. These appear as multiple coordinate triplets per residue, such as `(x_1,y_1,z_1)` and `(x_2,y_2,z_2)`.\n\n## 2.4 Submission table\n`sample_submission.csv` provides the exact required output layout:\n- You must output coordinates for five predicted structures.\n- You must keep the row order and the `ID` and `resname` fields intact.\n- `chain` and `copy` do not need to be provided.\n\nThis notebook uses `sample_submission.csv` as a template to avoid subtle formatting errors.\n","metadata":{}},{"id":"ae090cad","cell_type":"code","source":"# Load a small slice of each CSV and display schema\ndef read_csv_safe(path, nrows=None):\n    assert Path(path).exists(), f\"Missing file: {path}\"\n    return pd.read_csv(path, nrows=nrows)\n\ndf_train_seq = read_csv_safe(paths[\"train_sequences\"])\ndf_train_lbl = read_csv_safe(paths[\"train_labels\"])\ndf_val_seq   = read_csv_safe(paths[\"validation_sequences\"])\ndf_val_lbl   = read_csv_safe(paths[\"validation_labels\"])\ndf_test_seq  = read_csv_safe(paths[\"test_sequences\"])\ndf_sub       = read_csv_safe(paths[\"sample_submission\"])\n\nprint(\"train_sequences:\", df_train_seq.shape)\nprint(\"train_labels   :\", df_train_lbl.shape)\nprint(\"val_sequences  :\", df_val_seq.shape)\nprint(\"val_labels     :\", df_val_lbl.shape)\nprint(\"test_sequences :\", df_test_seq.shape)\nprint(\"sample_sub     :\", df_sub.shape)\n\ndisplay(df_train_seq.head(3))\ndisplay(df_train_lbl.head(3))\ndisplay(df_sub.head(3))\n\nprint(\"\\nSubmission columns:\")\nprint(df_sub.columns.tolist()[:20], \"...\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:46:18.049201Z","iopub.execute_input":"2026-01-09T18:46:18.049452Z","iopub.status.idle":"2026-01-09T18:46:25.377235Z","shell.execute_reply.started":"2026-01-09T18:46:18.049419Z","shell.execute_reply":"2026-01-09T18:46:25.376682Z"}},"outputs":[],"execution_count":null},{"id":"53bb05b9","cell_type":"markdown","source":"# 3. Dataset statistics and integrity checks\n\nA key operational step for any structure prediction workflow is to understand the distribution of sequence lengths and the presence of multi chain targets. For a target $t$ with sequence length $L_t$, the label table should contain exactly $L_t$ rows in the form `target_id_1` through `target_id_Lt`. Any deviation indicates either missing atoms, filtered residues, or an encoding inconsistency.\n\nWe will compute:\n1. sequence length distributions for each split\n2. label completeness for train and validation\n3. prevalence of multi conformation labels (multiple coordinate sets per residue)\n4. frequency of ligands and modified residues (via metadata)\n\nThese analyses are not only for exploration. They guide practical decisions:\n- maximum batch size for GPU inference\n- whether a single model can be used across the full length range\n- whether additional regularization is needed for long RNAs\n\nThe notebook also includes explicit assertions that protect against silent data leakage or misalignment between sequences and labels.\n","metadata":{}},{"id":"a7a83a50","cell_type":"code","source":"# Sequence length statistics\ndef add_len(df):\n    df = df.copy()\n    df[\"L\"] = df[\"sequence\"].astype(str).str.len()\n    return df\n\ntrain_len = add_len(df_train_seq)\nval_len   = add_len(df_val_seq)\ntest_len  = add_len(df_test_seq)\n\nstats = pd.DataFrame({\n    \"split\": [\"train\", \"validation\", \"test\"],\n    \"n_targets\": [len(train_len), len(val_len), len(test_len)],\n    \"L_min\": [train_len.L.min(), val_len.L.min(), test_len.L.min()],\n    \"L_med\": [train_len.L.median(), val_len.L.median(), test_len.L.median()],\n    \"L_mean\": [train_len.L.mean(), val_len.L.mean(), test_len.L.mean()],\n    \"L_max\": [train_len.L.max(), val_len.L.max(), test_len.L.max()],\n})\ndisplay(stats)\n\n# Label completeness check\ndef label_counts(df_lbl):\n    tmp = df_lbl.copy()\n    tmp[\"target_id\"] = tmp[\"ID\"].astype(str).apply(lambda x: \"_\".join(x.split(\"_\")[:-1]))\n    return tmp.groupby(\"target_id\")[\"resid\"].count().rename(\"n_rows\")\n\ntrain_counts = label_counts(df_train_lbl)\nval_counts   = label_counts(df_val_lbl)\n\ntrain_check = train_len.set_index(\"target_id\")[[\"L\"]].join(train_counts, how=\"left\")\nval_check   = val_len.set_index(\"target_id\")[[\"L\"]].join(val_counts, how=\"left\")\n\ntrain_check[\"missing_rows\"] = train_check[\"L\"] - train_check[\"n_rows\"]\nval_check[\"missing_rows\"]   = val_check[\"L\"] - val_check[\"n_rows\"]\n\nprint(\"Train targets with missing label rows:\", (train_check.missing_rows != 0).sum())\nprint(\"Val targets with missing label rows  :\", (val_check.missing_rows != 0).sum())\n\ndisplay(train_check.sort_values(\"missing_rows\", ascending=False).head(5))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:46:25.378122Z","iopub.execute_input":"2026-01-09T18:46:25.378438Z","iopub.status.idle":"2026-01-09T18:46:29.613274Z","shell.execute_reply.started":"2026-01-09T18:46:25.378411Z","shell.execute_reply":"2026-01-09T18:46:29.612633Z"}},"outputs":[],"execution_count":null},{"id":"4031269f","cell_type":"markdown","source":"# 4. The evaluation metric: TM score with residue indexed alignment\n\nThe competition score is the average across targets of the best TM score among your five submitted structures. For a single prediction and a single reference, TM score is:\n\n$$\n\\mathrm{TM}(\\hat{X}, X) = \\max_{R,\\,t}\\;\\frac{1}{L_{\\mathrm{ref}}}\n\\sum_{i=1}^{L_{\\mathrm{align}}}\n\\frac{1}{1 + \\left(\\frac{d_i(R,t)}{d_0}\\right)^2},\n$$\n\nwhere:\n- $R \\in SO(3)$ and $t \\in \\mathbb{R}^3$ are the rigid rotation and translation\n- $d_i(R,t) = \\lVert R\\hat{\\mathbf{x}}_i + t - \\mathbf{x}_i \\rVert_2$ is the distance after alignment\n- $L_{\\mathrm{ref}}$ is the number of residues in the reference\n- $L_{\\mathrm{align}}$ is the number of aligned residues\n\nThe scale parameter $d_0$ is defined as a function of $L_{\\mathrm{ref}}$. For longer chains:\n\n$$\nd_0(L) = \\sqrt[3]{L - 15} - 1.8, \\qquad L \\ge 30.\n$$\n\nFor short chains, the competition uses discrete constants:\n\n| $L_{\\mathrm{ref}}$ range | $d_0$ (Angstrom) |\n|---:|---:|\n| $L < 12$ | 0.3 |\n| 12 to 15 | 0.4 |\n| 16 to 19 | 0.5 |\n| 20 to 23 | 0.6 |\n| 24 to 29 | 0.7 |\n\nIn this competition, US align performs the alignment and it is configured to reward only residue pairs with matching residue numbers. Therefore, if your submission omits a residue or uses a different indexing convention, the metric cannot recover via alternative alignments.\n\nWe implement a local TM score routine (Kabsch alignment plus TM aggregation) for validation sanity checks. This local routine is not the authoritative scoring code, but it is useful for detecting pipeline regressions.\n","metadata":{}},{"id":"1b52c538","cell_type":"code","source":"# Kabsch alignment and TM-score implementation for sanity checks\ndef kabsch(P, Q, eps=1e-9):\n    \"\"\"Return R,t such that R@P + t best aligns to Q (least squares). P,Q: (L,3).\"\"\"\n    P = np.asarray(P, dtype=np.float64)\n    Q = np.asarray(Q, dtype=np.float64)\n    Pc = P - P.mean(axis=0, keepdims=True)\n    Qc = Q - Q.mean(axis=0, keepdims=True)\n    C = Pc.T @ Qc\n    V, S, Wt = np.linalg.svd(C)\n    d = np.sign(np.linalg.det(V @ Wt))\n    D = np.diag([1.0, 1.0, d])\n    R = V @ D @ Wt\n    t = Q.mean(axis=0) - (R @ P.mean(axis=0))\n    return R, t\n\ndef d0_from_L(Lref):\n    Lref = int(Lref)\n    if Lref >= 30:\n        return (Lref - 15) ** (1/3) - 1.8\n    # piecewise constants used in the competition description\n    if Lref < 12:\n        return 0.3\n    if 12 <= Lref <= 15:\n        return 0.4\n    if 16 <= Lref <= 19:\n        return 0.5\n    if 20 <= Lref <= 23:\n        return 0.6\n    return 0.7  # 24-29\n\ndef tm_score(P_pred, P_ref):\n    P_pred = np.asarray(P_pred, dtype=np.float64)\n    P_ref  = np.asarray(P_ref, dtype=np.float64)\n    assert P_pred.shape == P_ref.shape\n    Lref = P_ref.shape[0]\n    d0 = d0_from_L(Lref)\n    R, t = kabsch(P_pred, P_ref)\n    P_aln = (R @ P_pred.T).T + t\n    di = np.linalg.norm(P_aln - P_ref, axis=1)\n    return float(np.mean(1.0 / (1.0 + (di / d0) ** 2)))\n\n# quick sanity check: identical structures must give TM ~ 1\nexample_target = df_train_lbl[\"ID\"].iloc[0].split(\"_\")[0] + \"_\" + df_train_lbl[\"ID\"].iloc[0].split(\"_\")[1]\ntmp = df_train_lbl[df_train_lbl[\"ID\"].str.startswith(example_target + \"_\")].sort_values(\"resid\")\nP = tmp[[\"x_1\",\"y_1\",\"z_1\"]].to_numpy()\nprint(\"Example target:\", example_target, \"L=\", len(P))\nprint(\"TM(P,P) =\", tm_score(P, P))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:46:29.615143Z","iopub.execute_input":"2026-01-09T18:46:29.615461Z","iopub.status.idle":"2026-01-09T18:46:31.327291Z","shell.execute_reply.started":"2026-01-09T18:46:29.615439Z","shell.execute_reply":"2026-01-09T18:46:31.326453Z"}},"outputs":[],"execution_count":null},{"id":"4125f092","cell_type":"markdown","source":"# 5. Method overview\n\nBecause this is a code competition with a strict time budget, the highest leverage strategy is typically to start from a strong pretrained predictor and focus on:\n- correct data interfacing and residue indexing\n- robust conversion from model output to required CSV format\n- caching and runtime efficiency\n- generating multiple candidates (five structures) without excessive compute\n\nThe primary method in this notebook uses **RhoFold+** for inference. RhoFold+ is an open source system for RNA tertiary structure prediction. Its documented inference interface accepts:\n- an input FASTA file for the target sequence\n- optionally, a multiple sequence alignment file\n- a pretrained checkpoint path\n- a device selection (CPU or GPU)\n\nThe model produces a PDB file for an unrelaxed structure and optionally a relaxed structure after refinement. For this competition we only require the **C1'** atom coordinates, so we parse the PDB and extract that atom per residue. This approach remains fully compliant with the competition rules because it uses freely available pretrained models and does not require internet access at runtime.\n\n## Required attached datasets (Kaggle \"Add data\")\nTo make RhoFold+ available offline, attach:\n1. A dataset containing the RhoFold code repository.\n2. A dataset containing the pretrained weights file (checkpoint).\n\nThese datasets exist publicly on Kaggle (see the setup cell below). If they are missing, the notebook automatically falls back to a lightweight template based baseline that guarantees a valid submission.\n\n### Candidate generation strategy (five predictions)\nThe competition score takes the best of five TM scores per target. A computationally efficient strategy is to submit:\n- one high quality prediction (the main model output)\n- four additional candidates that are either identical or small perturbations of the main prediction\n\nIf the main prediction is correct, identical copies preserve its TM score under the best of five rule. If the model is stochastic, running multiple seeds can sometimes improve the best of five score but costs runtime. This notebook exposes both modes via a configuration flag.\n","metadata":{}},{"id":"2bc88c76","cell_type":"code","source":"# Configuration block\nclass CFG:\n    # Set to True for full test inference and CSV generation\n    MAKE_SUBMISSION = True\n\n    # Main model options\n    USE_RHOFOLD = True\n    RHOFOLD_SINGLE_SEQ = False  # use MSA when available (higher accuracy)\n    RHOFOLD_RELAX_STEPS = 0     # 0 uses unrelaxed structure\n\n    # MSA controls (RhoFold expects A3M; we generate multiple sampled A3M variants)\n    MSA_MAX_SEQS = 256          # cap MSA depth for speed/memory\n    RHOFOLD_RUNS_PER_TARGET = 3 # run RhoFold this many times; remaining predictions are duplicates\n    TRY_RELAX = False           # if True, attempt small Amber relaxation (may be slower / may fail)\n# Candidate generation\n    N_PRED = 5\n    STOCHASTIC_SAMPLES = True   # create multiple MSA variants / seeds for best-of-5\n\n    # Runtime controls\n    MAX_TARGETS = None          # set to an integer for debugging, otherwise None\n    VERBOSE = True\n\nprint(json.dumps(CFG.__dict__, indent=2, default=str))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:46:31.328274Z","iopub.execute_input":"2026-01-09T18:46:31.328558Z","iopub.status.idle":"2026-01-09T18:46:31.333826Z","shell.execute_reply.started":"2026-01-09T18:46:31.328535Z","shell.execute_reply":"2026-01-09T18:46:31.333105Z"}},"outputs":[],"execution_count":null},{"id":"d5259da9","cell_type":"markdown","source":"# 6. External assets and compliance with the offline runtime\n\nKaggle disables internet for final scoring. Therefore, any dependency that is not already installed in the base image must be provided as a local file. For RhoFold+ we use two inputs:\n- the RhoFold code itself\n- the pretrained checkpoint file\n\nThe official RhoFold repository describes both automatic download of the checkpoint and a manual download procedure. In the competition environment we cannot rely on automatic download, so we explicitly point to a local checkpoint path.\n\n## 6.1 Required attached datasets (Kaggle \"Add data\")\n\nAttach these public datasets to the notebook:\n1. **RhoFold code:** `andreashiropedi/rhofold` (contains `setup.py`, `inference.py`, and the `rhofold/` package).\n2. **RhoFold pretrained weights:** `shujun717/rhofold-pretrained-weights` (or any equivalent dataset that contains `rhofold_pretrained_params.pt` or `RhoFold_pretrained.pt`).\n\nIf your organization prefers a different packaging, you can attach an alternative dataset as long as it includes a compatible checkpoint and source tree. The detection logic in the next code cell searches for both the repository and a `.pt` checkpoint file.\n\n## 6.2 Offline installation strategy\n\nAfter locating the assets, the notebook:\n1. copies the repository to `/kaggle/working/rhofold_src/` (writeable)\n2. copies the checkpoint to `/kaggle/working/RhoFold_pretrained.pt`\n3. runs `python setup.py install --user` from the local source tree\n4. validates import with `import rhofold`\n\nIf any of these steps fail, the notebook falls back to the template baseline so that you can still submit.\n\nThis approach remains compliant with Kaggle rules: all external assets are freely available and attached through the platform, and no internet calls are made during execution.\n","metadata":{}},{"id":"4e0084c6","cell_type":"code","source":"# Locate and install RhoFold from attached Kaggle datasets (offline-safe)\n\nfrom pathlib import Path\nimport sys, os, glob, shutil\n\ndef _pick_rhofold_code_dir(input_root: Path) -> Path | None:\n    \"\"\"Return a directory that contains setup.py for RhoFold.\n\n    The Kaggle dataset `andreashiropedi/rhofold` is packaged with a top-level\n    folder named `RhoFold/` that contains `setup.py` and the `rhofold/` package.\n    In some mirrors, `setup.py` is at the dataset root. We support both layouts.\n    \"\"\"\n    candidates = [\n        input_root / \"rhofold\",\n        input_root / \"rhofold\" / \"RhoFold\",\n        input_root / \"RhoFold\",\n        input_root / \"RhoFold\" / \"RhoFold\",\n    ]\n    for c in candidates:\n        if (c / \"setup.py\").is_file():\n            return c\n\n    # Fallback: scan shallowly for setup.py to avoid expensive recursion\n    for c in [input_root / \"rhofold\", input_root / \"RhoFold\", input_root]:\n        if c.exists():\n            for p in c.glob(\"**/setup.py\"):\n                try:\n                    # Keep scan shallow-ish\n                    if len(p.parts) - len(c.parts) <= 3:\n                        return p.parent\n                except Exception:\n                    continue\n    return None\n\ndef _pick_rhofold_ckpt(input_root: Path) -> Path | None:\n    ckpt_candidates = [\n        input_root / \"rhofold-pretrained-weights\" / \"rhofold_pretrained.pt\",\n        input_root / \"rhofold_pretrained.pt\",\n        input_root / \"RhoFold_pretrained.pt\",\n    ]\n    for c in ckpt_candidates:\n        if c.is_file():\n            return c\n    # Fallback: search by name\n    for p in input_root.glob(\"**/*\"):\n        if p.is_file() and p.name.lower() in {\"rhofold_pretrained.pt\", \"rhofold-pretrained.pt\"}:\n            return p\n    return None\n\nrhofold_code_dir = _pick_rhofold_code_dir(INPUT_ROOT)\nrhofold_ckpt = _pick_rhofold_ckpt(INPUT_ROOT)\n\nprint(\"Detected rhofold_code_dir:\", rhofold_code_dir)\nprint(\"Detected rhofold_ckpt     :\", rhofold_ckpt)\n\nRHOFOLD_READY = False\nCKPT_PATH = None\nRHOFOLD_SRC = None\n\nif (rhofold_code_dir is not None) and (rhofold_ckpt is not None):\n    # Copy code to /kaggle/working so we can install/edit if needed\n    RHOFOLD_SRC = Path(\"/kaggle/working/rhofold_src\")\n    if RHOFOLD_SRC.exists():\n        shutil.rmtree(RHOFOLD_SRC)\n    RHOFOLD_SRC.mkdir(parents=True, exist_ok=True)\n\n    shutil.copytree(rhofold_code_dir, RHOFOLD_SRC / \"RhoFold\", dirs_exist_ok=True)\n\n    # Copy checkpoint into /kaggle/working for stable pathing\n    CKPT_PATH = Path(\"/kaggle/working/RhoFold_pretrained.pt\")\n    shutil.copy2(rhofold_ckpt, CKPT_PATH)\n\n    code_install_dir = RHOFOLD_SRC / \"RhoFold\"\n    print(\"Installed code dir :\", code_install_dir)\n    print(\"Checkpoint path    :\", CKPT_PATH)\n\n    # Offline-safe install: do not try to fetch dependencies from the internet\n    # Use editable mode for fast imports.\n    rc, _ = run([\"bash\",\"-lc\", f\"python -m pip install -q --user -e {code_install_dir} --no-deps\"])\nif rc != 0:\n    print(\"[WARN] Editable install failed (often harmless in offline Kaggle). Proceeding via sys.path only.\")\n\n    # Ensure import even if pip metadata is odd\n    sys.path.insert(0, str(code_install_dir))\n\n    try:\n        import rhofold  # noqa: F401\n        RHOFOLD_READY = True\n        print(\"Imported rhofold successfully.\")\n    except Exception as e:\n        RHOFOLD_READY = False\n        print(\"Failed to import rhofold:\", repr(e))\n        print(\"Tip: verify setup.py exists at\", code_install_dir / \"setup.py\")\n        print(\"Tip: list code dir:\"); run([\"bash\",\"-lc\", f\"ls -la {code_install_dir} | head -n 50\"])\n\nelse:\n    print(\"RhoFold assets not found. The notebook will use the fallback baseline.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:46:31.334772Z","iopub.execute_input":"2026-01-09T18:46:31.335062Z","iopub.status.idle":"2026-01-09T18:49:54.110901Z","shell.execute_reply.started":"2026-01-09T18:46:31.335038Z","shell.execute_reply":"2026-01-09T18:49:54.110034Z"}},"outputs":[],"execution_count":null},{"id":"81329f3e","cell_type":"markdown","source":"# 7. PDB parsing and extraction of C1' coordinates\n\nRhoFold+ produces PDB formatted coordinates. PDB is a fixed width text format where each atom is represented by an `ATOM` line containing:\n- atom name (columns 13 to 16)\n- residue name (nucleotide)\n- chain identifier\n- residue index\n- x, y, z coordinates in Angstroms\n\nFor this competition we must extract **only** the C1' atom per residue. The atom name is usually `C1'` in RNA PDB files. Some tools also output the equivalent `C1*`. The parser below accepts both.\n\nWe also enforce the coordinate clipping rule described in the dataset documentation. Coordinates are clipped to:\n$$\nx,y,z \\in [-999.999, 9999.999].\n$$\nThis is a legacy constraint derived from PDB field widths. If a model outputs extremely large coordinates, clipping prevents invalid formatting and matches the evaluation behavior.\n\nFinally, we return coordinates in the exact residue order expected by `sample_submission.csv`, which uses 1 based indexing.\n","metadata":{}},{"id":"72cf4b30","cell_type":"code","source":"# Robust PDB parser for C1' atoms (handles non-standard spacing)\nimport re\nimport numpy as np\n\ndef parse_pdb_c1_robust(pdb_path: str, L_expected: int | None = None) -> np.ndarray:\n    \"\"\"Parse C1' (or C1*) coordinates from a PDB file robustly.\n\n    Why this exists:\n    - Some ML-generated PDBs (and even hand-written quick tests) do not strictly\n      respect fixed-column alignment. Fixed slicing like line[22:26] can fail.\n    - We therefore prefer whitespace splitting, while still accepting C1' / C1*.\n\n    Returns\n    -------\n    X : (L,3) float32 ndarray\n        Ordered by 1-based residue index (resid). Missing residues are filled\n        by nearest-neighbor copying (then clipped to the competition's legacy range).\n    \"\"\"\n    coords: dict[int, tuple[float,float,float]] = {}\n\n    def _to_int_resid(tok: str):\n        # Accept tokens like '12', '12A', '-3', '0042'\n        m = re.match(r\"^(-?\\d+)\", tok)\n        return int(m.group(1)) if m else None\n\n    with open(pdb_path, \"r\") as f:\n        for line in f:\n            if not (line.startswith(\"ATOM\") or line.startswith(\"HETATM\")):\n                continue\n            parts = line.split()\n            if len(parts) < 8:\n                continue\n\n            atom = parts[2]\n            if atom not in (\"C1'\", \"C1*\"):\n                continue\n\n            # Typical layouts:\n            # ATOM serial atom resname chain resid x y z ...\n            # ATOM serial atom resname resid x y z ...  (chain omitted)\n            if _to_int_resid(parts[4]) is not None:\n                resid_tok = parts[4]\n                xyz_start = 5\n            else:\n                if len(parts) < 9:\n                    continue\n                resid_tok = parts[5]\n                xyz_start = 6\n\n            resid = _to_int_resid(resid_tok)\n            if resid is None:\n                continue\n\n            try:\n                x = float(parts[xyz_start + 0])\n                y = float(parts[xyz_start + 1])\n                z = float(parts[xyz_start + 2])\n            except Exception:\n                continue\n\n            if resid not in coords:\n                coords[resid] = (x, y, z)\n\n    if not coords:\n        raise ValueError(f\"No C1' (or C1*) atoms parsed from: {pdb_path}\")\n\n    max_resid = max(coords.keys())\n    L = int(L_expected) if L_expected is not None else max_resid\n\n    X = np.zeros((L, 3), dtype=np.float32)\n    missing = []\n\n    for i in range(1, L + 1):\n        if i in coords:\n            X[i - 1] = coords[i]\n        else:\n            missing.append(i)\n\n    # Fill missing residues by nearest neighbor to avoid crashing downstream\n    if missing:\n        for i in missing:\n            left = i - 1\n            right = i + 1\n            while left >= 1 and left not in coords:\n                left -= 1\n            while right <= L and right not in coords:\n                right += 1\n            if left >= 1 and left in coords:\n                X[i - 1] = coords[left]\n            elif right <= L and right in coords:\n                X[i - 1] = coords[right]\n            else:\n                X[i - 1] = 0.0\n        print(f\"[WARN] Missing C1' for {len(missing)} residues in {pdb_path}. Filled by neighbors/zeros.\")\n\n    # Competition uses legacy 8-char coordinate field; clamp to scoring range\n    X = np.clip(X, -999.999, 9999.999).astype(np.float32)\n    # Sanitize NaN/Inf just in case the PDB contains 'nan' values\n    X = np.nan_to_num(X, nan=0.0, posinf=9999.999, neginf=-999.999).astype(np.float32)\n    return X\n# Backwards-compatible alias used elsewhere in the notebook\ndef parse_pdb_c1(pdb_path, L_expected=None):\n    return parse_pdb_c1_robust(str(pdb_path), L_expected=L_expected)\n\n# quick self-test with a synthetic minimal PDB (non-standard spacing is intentional)\npdb_test = WORKDIR / \"pdb_test.pdb\"\npdb_test.write_text(\n    \"ATOM      1  C1'  G   A   1       1.000   2.000   3.000  1.00 20.00           C\\n\"\n    \"ATOM      2  C1'  A   A   2       4.000   5.000   6.000  1.00 20.00           C\\n\"\n    \"END\\n\"\n)\nX = parse_pdb_c1(pdb_test, L_expected=2)\nprint(\"Parsed coords shape:\", X.shape)\nprint(X)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:49:54.111899Z","iopub.execute_input":"2026-01-09T18:49:54.112147Z","iopub.status.idle":"2026-01-09T18:49:54.126237Z","shell.execute_reply.started":"2026-01-09T18:49:54.112125Z","shell.execute_reply":"2026-01-09T18:49:54.125508Z"}},"outputs":[],"execution_count":null},{"id":"2c73d0cf","cell_type":"markdown","source":"# 8. RhoFold+ inference wrapper\n\nThe RhoFold repository exposes an inference script with command line arguments. According to the documentation, the key flags are:\n- `--input_fas`: path to FASTA file containing the RNA sequence\n- `--input_a3m`: optional MSA file\n- `--output_dir`: directory where PDB and auxiliary outputs are written\n- `--device`: `cpu` or `cuda:<idx>`\n- `--ckpt`: path to pretrained checkpoint\n- `--relax_steps`: number of refinement steps\n\nWe wrap this interface into a Python function that:\n1. writes a target specific FASTA file\n2. optionally converts the provided MSA FASTA to a permissive A3M like format\n3. calls the inference script as a subprocess (robust across package layouts)\n4. selects either `unrelaxed_model.pdb` or `relaxed_{steps}_model.pdb` depending on configuration\n5. extracts C1' coordinates using the parser from the previous section\n\nTo keep runtime under control, this wrapper caches results in `/kaggle/working/cache_rhofold/`. If you rerun the notebook, previously predicted structures are reused.\n","metadata":{}},{"id":"828d5290","cell_type":"code","source":"import hashlib\nimport re\nimport numpy as np\nfrom pathlib import Path\n# RhoFold inference wrapper (subprocess based)\n\nCACHE_DIR = WORKDIR / \"cache_rhofold\"\nCACHE_DIR.mkdir(parents=True, exist_ok=True)\n\ndef write_fasta(target_id, seq, out_dir):\n    out_dir = Path(out_dir)\n    out_dir.mkdir(parents=True, exist_ok=True)\n    fas = out_dir / f\"{target_id}.fasta\"\n    fas.write_text(f\">{target_id}\\n{seq}\\n\")\n    return fas\n\ndef _read_fasta_records(path):\n    records = []\n    with open(path, \"r\") as f:\n        header, buf = None, []\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                if header is not None:\n                    records.append((header, \"\".join(buf)))\n                header = line[1:].strip()\n                buf = []\n            else:\n                buf.append(line)\n        if header is not None:\n            records.append((header, \"\".join(buf)))\n    return records\n\ndef _sanitize_msa_seq(s: str) -> str:\n    # Keep A,C,G,U and gaps; everything else becomes gap.\n    return re.sub(r\"[^ACGU\\-]\", \"-\", s.upper())\n\ndef write_a3m_variants(target_id, msa_fasta_path, out_dir, n_variants=5, max_seqs=256, seed_base=0):\n    \"\"\"\n    Convert Kaggle-provided MSA FASTA to multiple A3M variants by random subsampling.\n\n    RhoFold supports `--input_a3m` and is notably more accurate with MSA than in single-seq mode.\n    We generate multiple sampled A3Ms so the competition's best-of-5 scoring can pick the best conformation.\n    \"\"\"\n    if msa_fasta_path is None or (not Path(msa_fasta_path).exists()):\n        return [None] * n_variants\n\n    out_dir = Path(out_dir)\n    out_dir.mkdir(parents=True, exist_ok=True)\n\n    records = _read_fasta_records(msa_fasta_path)\n    if not records:\n        return [None] * n_variants\n\n    # Use the first sequence as the query (common convention for MSAs)\n    qh, qs = records[0]\n    query = (f\"{target_id}|query\", _sanitize_msa_seq(qs))\n    others = [(h, _sanitize_msa_seq(s)) for (h, s) in records[1:]]\n\n    variants = []\n    for v in range(n_variants):\n        rng = np.random.default_rng(seed_base + v)\n        if len(others) <= (max_seqs - 1):\n            picked = others.copy()\n            rng.shuffle(picked)\n        else:\n            idx = rng.choice(len(others), size=(max_seqs - 1), replace=False)\n            picked = [others[i] for i in idx]\n        a3m = out_dir / f\"{target_id}.v{v}.a3m\"\n        with open(a3m, \"w\") as g:\n            g.write(f\">{query[0]}\\n{query[1]}\\n\")\n            for h, s in picked:\n                g.write(f\">{h}\\n{s}\\n\")\n        variants.append(a3m)\n    return variants\n\ndef rhofold_predict_coords(target_id, seq, msa_path=None, a3m_path=None, device=\"cuda:0\", relax_steps=None):\n    assert RHOFOLD_READY, \"RhoFold is not installed or checkpoint missing.\"\n    # Cache key\n    key = hashlib.md5((target_id + seq + str(CFG.RHOFOLD_SINGLE_SEQ) + str(relax_steps if relax_steps is not None else CFG.RHOFOLD_RELAX_STEPS) + str(a3m_path) ).encode()).hexdigest()[:16]\n    out_dir = CACHE_DIR / f\"{target_id}_{key}\"\n    pdb_out = out_dir / \"unrelaxed_model.pdb\"\n    rs = int(relax_steps if relax_steps is not None else CFG.RHOFOLD_RELAX_STEPS)\n    if rs and rs > 0:\n        pdb_out = out_dir / f\"relaxed_{rs}_model.pdb\"\n    if pdb_out.exists():\n        return parse_pdb_c1(pdb_out, L_expected=len(seq))\n\n    fas = write_fasta(target_id, seq, out_dir)\n    a3m = None\n    if a3m_path is not None:\n        a3m = Path(a3m_path)\n    elif (not CFG.RHOFOLD_SINGLE_SEQ) and (msa_path is not None):\n        # Backward-compatible: single A3M derived from full MSA\n        a3m = write_a3m_variants(target_id, msa_path, out_dir, n_variants=1, max_seqs=CFG.MSA_MAX_SEQS, seed_base=0)[0]\n\n    # find inference.py location inside installed source copy\n    inference_candidates = list((WORKDIR/\"rhofold_src\").rglob(\"inference.py\"))\n    if not inference_candidates:\n        raise FileNotFoundError(\"Could not locate inference.py in rhofold_src\")\n    inference_py = str(inference_candidates[0])\n\n    ckpt = str((WORKDIR/\"RhoFold_pretrained.pt\").resolve())\n    cmd = [\"python\", inference_py,\n           \"--input_fas\", str(fas),\n           \"--output_dir\", str(out_dir),\n           \"--ckpt\", ckpt,\n           \"--device\", device,\n           \"--relax_steps\", str(int(relax_steps if relax_steps is not None else CFG.RHOFOLD_RELAX_STEPS))]\n    if CFG.RHOFOLD_SINGLE_SEQ:\n        cmd += [\"--single_seq_pred\", \"True\"]\n    else:\n        if a3m is not None:\n            cmd += [\"--input_a3m\", str(a3m)]\n    rc, out = run([\"bash\",\"-lc\", \" \".join(cmd)])\n    if rc != 0:\n        raise RuntimeError(f\"RhoFold inference failed for {target_id}\")\n\n    if not pdb_out.exists():\n        # Fall back to unrelaxed if relaxed not produced\n        pdb_out = out_dir / \"unrelaxed_model.pdb\"\n    return parse_pdb_c1(pdb_out, L_expected=len(seq))\n\nprint(\"RhoFold wrapper ready:\", RHOFOLD_READY)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:49:54.127113Z","iopub.execute_input":"2026-01-09T18:49:54.127602Z","iopub.status.idle":"2026-01-09T18:49:54.149745Z","shell.execute_reply.started":"2026-01-09T18:49:54.127567Z","shell.execute_reply":"2026-01-09T18:49:54.14913Z"}},"outputs":[],"execution_count":null},{"id":"bb8e6933","cell_type":"markdown","source":"# 9. Fallback baseline: nearest neighbor template transfer\n\nA robust competition notebook should never fail to generate a submission. External assets can be missing, the accelerator may be disabled, or the inference code can fail for a rare long sequence. To guarantee a valid output, we include a fallback baseline that is purely data driven and runs on CPU.\n\nThe baseline builds a library of training structures and attempts to find a close sequence match for each test target. If a suitable match exists, we reuse the training coordinates as a template. Otherwise we generate a simple A form like helix as a placeholder.\n\nThis baseline is not intended to win the leaderboard, but it ensures:\n- correct file format\n- correct residue indexing\n- deterministic outputs\n\nThe Part 1 outcome preprint notes that top performing strategies can be template based in many settings, so even a lightweight nearest neighbor template can provide a non trivial score on homologous targets.\n","metadata":{}},{"id":"332adf09","cell_type":"code","source":"# Build a minimal template library from train_labels for fallback\n\ndef build_train_library(df_seq, df_lbl):\n    # Map target_id -> (sequence, coords)\n    lib = {}\n    df_lbl = df_lbl.copy()\n    df_lbl[\"target_id\"] = df_lbl[\"ID\"].apply(lambda x: \"_\".join(x.split(\"_\")[:-1]))\n    for tid, grp in df_lbl.groupby(\"target_id\"):\n        grp = grp.sort_values(\"resid\")\n        coords = grp[[\"x_1\",\"y_1\",\"z_1\"]].to_numpy(dtype=np.float32)\n        seq = df_seq.loc[df_seq.target_id == tid, \"sequence\"]\n        if len(seq) == 0:\n            continue\n        seq = str(seq.iloc[0])\n        if len(seq) != len(coords):\n            # allow mismatch, truncate to min length\n            L = min(len(seq), len(coords))\n            seq = seq[:L]\n            coords = coords[:L]\n        lib[tid] = (seq, coords)\n    return lib\n\nTRAIN_LIB = build_train_library(df_train_seq, df_train_lbl)\nprint(\"Template library size:\", len(TRAIN_LIB))\n\ndef jaccard_kmer(a, b, k=5):\n    if len(a) < k or len(b) < k:\n        return 0.0\n    sa = set(a[i:i+k] for i in range(len(a)-k+1))\n    sb = set(b[i:i+k] for i in range(len(b)-k+1))\n    inter = len(sa & sb)\n    union = len(sa | sb) + 1e-9\n    return inter / union\n\ndef find_best_template(seq, lib, k=5):\n    best = (None, -1.0)\n    for tid, (s, coords) in lib.items():\n        sim = jaccard_kmer(seq, s, k=k)\n        if sim > best[1]:\n            best = (tid, sim)\n    return best\n\ndef helix_coords(L, rise=3.4, twist=32.7 * math.pi/180.0, radius=8.0):\n    # Simple coarse helical backbone around z-axis\n    X = np.zeros((L,3), dtype=np.float32)\n    for i in range(L):\n        ang = i * twist\n        X[i,0] = radius * math.cos(ang)\n        X[i,1] = radius * math.sin(ang)\n        X[i,2] = i * rise\n    return X\n\ndef fallback_predict(seq):\n    tid, sim = find_best_template(seq, TRAIN_LIB)\n    if tid is not None and sim > 0.35:  # heuristic threshold\n        tpl_seq, tpl_coords = TRAIN_LIB[tid]\n        L = min(len(seq), tpl_coords.shape[0])\n        X = tpl_coords[:L].copy()\n        if len(seq) > L:\n            X2 = helix_coords(len(seq)-L) + X[-1]\n            X = np.vstack([X, X2])\n        return X\n    return helix_coords(len(seq))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:49:54.150636Z","iopub.execute_input":"2026-01-09T18:49:54.150887Z","iopub.status.idle":"2026-01-09T18:50:06.981731Z","shell.execute_reply.started":"2026-01-09T18:49:54.150857Z","shell.execute_reply":"2026-01-09T18:50:06.98098Z"}},"outputs":[],"execution_count":null},{"id":"2ed99bc4","cell_type":"markdown","source":"# 10. End to end inference for the test set\n\nWe now assemble the complete inference loop that reads `test_sequences.csv` and produces a `submission.csv` matching `sample_submission.csv`.\n\n## 10.1 Target iteration and caching\nFor each target:\n1. read `sequence`\n2. locate the corresponding MSA file if needed: `MSA/{target_id}.MSA.fasta`\n3. predict a coordinate array $\\hat{X}\\in\\mathbb{R}^{L\\times 3}$\n4. map the array to the submission rows for that target\n\nIf RhoFold is available and enabled, we run RhoFold inference. If it is unavailable or fails for a target, we fall back to template transfer.\n\n## 10.2 Mapping predicted coordinates to the submission dataframe\n`sample_submission.csv` includes rows for each residue, already labeled with `ID`, `resname`, and `resid`. We parse `target_id` from `ID` by dropping the final underscore and resid component. Then we fill:\n\n$$\nx_k[i] = \\hat{x}_i, \\quad y_k[i] = \\hat{y}_i, \\quad z_k[i] = \\hat{z}_i, \\quad k=1,\\dots,5.\n$$\n\nIf `STOCHASTIC_SAMPLES = False`, the same $\\hat{X}$ is reused for all five predictions. This is an efficient default under a strict runtime budget.\n","metadata":{}},{"id":"f7b205c7","cell_type":"code","source":"# End to end inference on test targets and fill sample submission (robust)\n\nimport numpy as np\nimport pandas as pd\nimport time\nimport hashlib  # safety: used by some caching helpers elsewhere\n\n# Start from sample_submission to preserve row order\nsub = df_sub.copy()\n\n# Add helper columns\nsub[\"target_id\"] = sub[\"ID\"].astype(str).apply(lambda x: \"_\".join(x.split(\"_\")[:-1]))\nsub[\"resid\"] = sub[\"resid\"].astype(int)\n\n# Coordinate columns (5 models)\ncoord_cols = [f\"{ax}_{k}\" for k in range(1, CFG.N_PRED + 1) for ax in [\"x\", \"y\", \"z\"]]\n\n# Ensure all coord cols exist and are float32\nfor col in coord_cols:\n    if col not in sub.columns:\n        sub[col] = 0.0\nsub[coord_cols] = sub[coord_cols].astype(np.float32)\nsub[coord_cols] = sub[coord_cols].replace([np.inf, -np.inf], np.nan).fillna(0.0)\n\n# Device selection (prefer GPU if available)\ndevice = \"cpu\"\ntry:\n    import torch\n    if torch.cuda.is_available():\n        device = \"cuda:0\"\nexcept Exception:\n    pass\nprint(\"Selected device:\", device)\n\n# Target list: use official test_sequences.csv by default\ntarget_df = df_test_seq if \"df_test_seq\" in globals() else df_val_seq\nassert \"target_id\" in target_df.columns and \"sequence\" in target_df.columns, \"Expected columns: target_id, sequence\"\ntarget_ids = target_df[\"target_id\"].astype(str).tolist()\ntarget_seqs = target_df[\"sequence\"].astype(str).tolist()\nprint(\"Targets to predict:\", len(target_ids))\n\npred_cache = {}   # tid -> list[np.ndarray] length N_PRED, each (L,3)\nfailures = []\nt0 = time.time()\n\nfor i, (tid, seq) in enumerate(zip(target_ids, target_seqs), start=1):\n    try:\n        L = len(seq)\n        msa_path = MSA_DIR / f\"{tid}.MSA.fasta\"\n\n        # Prepare A3M variants for diversity (best-of-5), if available\n        a3m_variants = [None] * CFG.N_PRED\n        if (\n            CFG.USE_RHOFOLD\n            and RHOFOLD_READY\n            and (not CFG.RHOFOLD_SINGLE_SEQ)\n            and msa_path.exists()\n        ):\n            a3m_variants = write_a3m_variants(\n                tid=tid,\n                msa_fasta_path=str(msa_path),\n                out_dir=(CACHE_DIR / f\"{tid}_msa\"),\n                n_variants=CFG.N_PRED,\n                max_seqs=CFG.MSA_MAX_SEQS,\n                seed_base=abs(hash(tid)) % (2**31),\n            )\n\n        # Run RhoFold only a few times to meet time budget; duplicate remaining predictions\n        n_runs = min(int(CFG.RHOFOLD_RUNS_PER_TARGET), CFG.N_PRED)\n        X_list = []\n\n        for k in range(n_runs):\n            if CFG.USE_RHOFOLD and RHOFOLD_READY:\n                relax_k = 0\n                if CFG.TRY_RELAX and k == 0:\n                    relax_k = 50\n\n                Xk = rhofold_predict_coords(\n                    tid=tid,\n                    seq=seq,\n                    msa_path=str(msa_path) if msa_path.exists() else None,\n                    a3m_path=str(a3m_variants[k]) if a3m_variants[k] is not None else None,\n                    device=device,\n                    relax_steps=relax_k,\n                )\n            else:\n                Xk = fallback_predict(seq)\n\n            Xk = np.asarray(Xk, dtype=np.float32)\n            Xk = np.nan_to_num(Xk, nan=0.0, posinf=9999.999, neginf=-999.999)\n            if Xk.shape[0] != L:\n                # strict safety: force correct length\n                Xk_fixed = np.zeros((L, 3), dtype=np.float32)\n                m = min(L, Xk.shape[0])\n                Xk_fixed[:m] = Xk[:m]\n                Xk = Xk_fixed\n\n            X_list.append(Xk)\n\n        # If nothing produced, fallback once\n        if len(X_list) == 0:\n            X0 = np.asarray(fallback_predict(seq), dtype=np.float32)\n            X0 = np.nan_to_num(X0, nan=0.0, posinf=9999.999, neginf=-999.999)\n            if X0.shape[0] != L:\n                X0_fixed = np.zeros((L, 3), dtype=np.float32)\n                m = min(L, X0.shape[0])\n                X0_fixed[:m] = X0[:m]\n                X0 = X0_fixed\n            X_list = [X0]\n\n        # Duplicate last prediction to reach N_PRED\n        while len(X_list) < CFG.N_PRED:\n            X_list.append(X_list[-1].copy())\n\n        pred_cache[tid] = X_list\n\n    except Exception as e:\n        failures.append((tid, repr(e)))\n        L = len(seq)\n        X0 = np.asarray(fallback_predict(seq), dtype=np.float32)\n        X0 = np.nan_to_num(X0, nan=0.0, posinf=9999.999, neginf=-999.999)\n        if X0.shape[0] != L:\n            X0_fixed = np.zeros((L, 3), dtype=np.float32)\n            m = min(L, X0.shape[0])\n            X0_fixed[:m] = X0[:m]\n            X0 = X0_fixed\n        pred_cache[tid] = [X0.copy() for _ in range(CFG.N_PRED)]\n\n    if CFG.VERBOSE and (i % 10 == 0 or i == 1):\n        print(f\"[{i:4d}/{len(target_ids)}] done {tid}  L={len(seq)}  elapsed={time.time()-t0:7.1f}s\")\n\nprint(\"Targets processed:\", len(target_ids), \"Failures:\", len(failures))\nif failures:\n    print(\"Failure examples:\", failures[:3])\n\n# Fill submission\nfor tid, X_list in pred_cache.items():\n    mask = (sub[\"target_id\"] == tid)\n    if not mask.any():\n        continue\n\n    res = sub.loc[mask, \"resid\"].to_numpy(dtype=np.int32) - 1  # 0-based\n    nrows = res.shape[0]\n\n    for k in range(1, CFG.N_PRED + 1):\n        Xk = X_list[k - 1]\n        Xk = np.asarray(Xk, dtype=np.float32)\n        Xk = np.nan_to_num(Xk, nan=0.0, posinf=9999.999, neginf=-999.999)\n\n        good = (res >= 0) & (res < Xk.shape[0])\n        rows_xyz = np.zeros((nrows, 3), dtype=np.float32)\n        if np.any(good):\n            rows_xyz[good] = Xk[res[good]]\n\n        sub.loc[mask, f\"x_{k}\"] = rows_xyz[:, 0]\n        sub.loc[mask, f\"y_{k}\"] = rows_xyz[:, 1]\n        sub.loc[mask, f\"z_{k}\"] = rows_xyz[:, 2]\n\n# Final sanitize + clip\nsub[coord_cols] = sub[coord_cols].astype(np.float32)\nsub[coord_cols] = sub[coord_cols].replace([np.inf, -np.inf], np.nan).fillna(0.0)\nsub[coord_cols] = sub[coord_cols].clip(-999.999, 9999.999)\n\ndisplay(sub.head(3))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:56:52.312855Z","iopub.execute_input":"2026-01-09T18:56:52.313768Z","iopub.status.idle":"2026-01-09T18:57:36.862782Z","shell.execute_reply.started":"2026-01-09T18:56:52.313729Z","shell.execute_reply":"2026-01-09T18:57:36.862113Z"}},"outputs":[],"execution_count":null},{"id":"7b3c0e95","cell_type":"markdown","source":"# 11. Format validation and submission writing\n\nEven high quality predictions will score zero if the submission file violates formatting requirements. We therefore validate:\n\n1. The submission has exactly the same number of rows as `sample_submission.csv`.\n2. The `ID` column matches the template exactly.\n3. All required coordinate columns exist: `x_1` through `z_5`.\n4. All coordinate values are finite and within the allowed clipping range.\n\nOnce validated, we write `submission.csv` to the working directory. Kaggle will pick this file up for scoring.\n","metadata":{}},{"id":"747b94e8","cell_type":"code","source":"# Validate submission integrity and write to disk\n\nimport numpy as np\n\nrequired_cols = [\"ID\", \"resname\", \"resid\"] + [f\"{ax}_{k}\" for k in range(1,6) for ax in [\"x\",\"y\",\"z\"]]\nmissing = [c for c in required_cols if c not in sub.columns]\nassert not missing, f\"Missing columns: {missing}\"\n\nassert len(sub) == len(df_sub), \"Row count mismatch with sample_submission\"\nassert (sub[\"ID\"].values == df_sub[\"ID\"].values).all(), \"ID mismatch vs sample_submission\"\n\ncoord_cols = [c for c in sub.columns if c.startswith((\"x_\",\"y_\",\"z_\"))]\nsub[coord_cols] = sub[coord_cols].astype(np.float32)\nsub[coord_cols] = sub[coord_cols].replace([np.inf, -np.inf], np.nan).fillna(0.0)\nassert np.isfinite(sub[coord_cols].to_numpy()).all(), \"Non-finite coordinates detected\"\n\n# Clip exactly to legacy limits (as described in competition)\nsub[coord_cols] = sub[coord_cols].clip(-999.999, 9999.999)\n\nout_path = WORKDIR / \"submission.csv\"\nsub[required_cols].to_csv(out_path, index=False)\nprint(\"Wrote:\", out_path, \"size=\", out_path.stat().st_size)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:58:05.946204Z","iopub.execute_input":"2026-01-09T18:58:05.946976Z","iopub.status.idle":"2026-01-09T18:58:06.096142Z","shell.execute_reply.started":"2026-01-09T18:58:05.946945Z","shell.execute_reply":"2026-01-09T18:58:06.095552Z"}},"outputs":[],"execution_count":null},{"id":"33f6c4b9","cell_type":"markdown","source":"# 12. Optional local validation using the public validation split\n\nThe public validation split provides labeled coordinates that are similar in nature to the hidden test. While the official scoring uses US align and may include multiple reference conformations, a local evaluation can still be useful for:\n- confirming that the inference pipeline is functional\n- comparing alternative configuration choices (single sequence prediction versus MSA)\n- detecting errors in residue indexing\n\nWe compute a local TM score using the `tm_score` routine defined earlier. The evaluation uses the first label structure `(x_1,y_1,z_1)` as a single reference. This is a simplification of the official procedure, but it offers a stable signal for debugging.\n\nThis block is optional and can be disabled if you are running near the time limit.\n","metadata":{}},{"id":"82a923f1","cell_type":"code","source":"# Local validation on a subset (optional)\n# DO_LOCAL_VAL = True and (CFG.MAX_TARGETS is not None or len(df_val_seq) <= 200)\n\n# if DO_LOCAL_VAL:\n#     val_map = dict(zip(df_val_seq[\"target_id\"], df_val_seq[\"sequence\"]))\n#     val_targets = list(val_map.keys())[: min(50, len(val_map))]\n#     scores = []\n#     for tid in val_targets:\n#         seq = str(val_map[tid])\n#         msa_path = paths[\"MSA_dir\"] / f\"{tid}.MSA.fasta\"\n#         # reference\n#         ref = df_val_lbl[df_val_lbl[\"ID\"].str.startswith(tid + \"_\")].sort_values(\"resid\")[[\"x_1\",\"y_1\",\"z_1\"]].to_numpy()\n#         try:\n#             if CFG.USE_RHOFOLD and RHOFOLD_READY:\n#                 pred = rhofold_predict_coords(tid, seq, msa_path=str(msa_path) if msa_path.exists() else None, device=device)\n#             else:\n#                 pred = fallback_predict(seq)\n#             L = min(len(ref), pred.shape[0])\n#             score = tm_score(pred[:L], ref[:L])\n#         except Exception:\n#             score = 0.0\n#         scores.append(score)\n#     print(\"Local TM (subset) mean:\", float(np.mean(scores)), \"median:\", float(np.median(scores)))\n# else:\n#     print(\"Skipping local validation.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:58:48.930442Z","iopub.execute_input":"2026-01-09T18:58:48.931173Z","iopub.status.idle":"2026-01-09T18:58:48.934611Z","shell.execute_reply.started":"2026-01-09T18:58:48.931144Z","shell.execute_reply":"2026-01-09T18:58:48.933938Z"}},"outputs":[],"execution_count":null},{"id":"152df072","cell_type":"markdown","source":"# 13. Computational considerations and runtime budgeting\n\n## 13.1 Accelerator choice\nThis notebook targets **GPU P100 (16 GB)**, a commonly available Kaggle accelerator. GPU acceleration is beneficial because modern structure predictors are dominated by dense tensor operations and attention like blocks. Kaggle can also provide T4 class GPUs and TPU v3-8 in some environments, but using TPU typically requires a different software stack and introduces additional engineering complexity.\n\nA practical comparison is:\n\n| Accelerator | Typical memory | Strengths for this notebook | Notes |\n|---|---:|---|---|\n| NVIDIA Tesla P100 | 16 GB | Strong FP32 throughput, simple CUDA workflow | Good default in Kaggle notebooks |\n| NVIDIA T4 (sometimes x2) | 16 GB per GPU | Tensor cores for FP16, good inference | Multi GPU adds engineering overhead |\n| TPU v3-8 | 8 cores, about 16 GB per core | Very high throughput for XLA stacks | Requires JAX or TF plus data pipeline changes |\n\nThe notebook also introspects the actual hardware at runtime using `nvidia-smi` and `lscpu`.\n\n## 13.2 Time complexity drivers\nLet $L$ be sequence length and $M$ the number of MSA sequences used. Many modern predictors have components with quadratic complexity:\n- self attention: $\\mathcal{O}(L^2 d)$\n- pair representation updates: $\\mathcal{O}(L^2)$\n- distogram or contact logits: $\\mathcal{O}(L^2)$\n\nFor long targets, these can dominate runtime and memory. To keep inference under a strict 1 hour budget, we:\n- default to `single_seq_pred` mode for RhoFold when available\n- disable expensive refinement steps (`relax_steps = 0`)\n- cache all intermediate predictions\n- avoid running five independent stochastic samples by default\n\n## 13.3 Memory budgeting\nThe dominant memory terms in attention based architectures scale as:\n$$\n\\mathcal{O}(L^2) \\text{ for pair features and attention maps, and } \\mathcal{O}(L d) \\text{ for sequence embeddings}.\n$$\nIf you observe out of memory errors for long sequences, reduce the batch size, enable mixed precision if supported by the underlying model, or split the sequence into domains with a downstream stitching procedure.\n\n## 13.4 Reproducibility\nWe fix NumPy and Python random seeds. RhoFold inference is generally deterministic in single sequence mode. If you enable stochastic sampling, additional seeds must be set for PyTorch as well.\n","metadata":{}},{"id":"0b2ea537","cell_type":"code","source":"# Optional: set PyTorch deterministic settings when stochastic sampling is enabled\ntry:\n    import torch\n    torch.manual_seed(SEED)\n    torch.cuda.manual_seed_all(SEED)\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    print(\"PyTorch deterministic flags set.\")\nexcept Exception as e:\n    print(\"PyTorch not available or could not set flags:\", repr(e))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:58:52.851683Z","iopub.execute_input":"2026-01-09T18:58:52.851979Z","iopub.status.idle":"2026-01-09T18:58:52.861867Z","shell.execute_reply.started":"2026-01-09T18:58:52.851954Z","shell.execute_reply":"2026-01-09T18:58:52.861198Z"}},"outputs":[],"execution_count":null},{"id":"a163f683","cell_type":"markdown","source":"# 14. Implementation details for candidate generation (five structures)\n\nThe competition requires five coordinate sets per target, and scoring uses the best TM score among the five. We expose two operational modes:\n\n## Mode A (default): replicate\nLet $\\hat{X}$ be the main predicted structure. Set:\n$$\n\\hat{X}^{(1)} = \\hat{X}^{(2)} = \\dots = \\hat{X}^{(5)} = \\hat{X}.\n$$\nThis yields best of five identical to the main prediction and has minimal runtime.\n\n## Mode B: stochastic sampling\nRun the model multiple times with different random seeds, stochastic dropout, or different MSA sampling. This can improve the best of five score if the model produces diverse plausible folds. The cost is roughly a factor of five in runtime.\n\nIn this notebook, Mode B is disabled by default to satisfy the strict runtime requirement. If you want to experiment with Mode B, set `CFG.STOCHASTIC_SAMPLES = True` and update the inference loop to call `rhofold_predict_coords` multiple times per target.\n","metadata":{}},{"id":"6dc1fe3d","cell_type":"code","source":"# Placeholder for advanced candidate generation\n# For now, we use replicated coordinates. To implement stochastic sampling, you can:\n# 1) Add a loop over seeds\n# 2) Toggle dropout inside the model if supported\n# 3) Use alternative MSAs (clustered, sampled) if available\n\nprint(\"STOCHASTIC_SAMPLES =\", CFG.STOCHASTIC_SAMPLES, \"N_PRED =\", CFG.N_PRED)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:58:59.508983Z","iopub.execute_input":"2026-01-09T18:58:59.50929Z","iopub.status.idle":"2026-01-09T18:58:59.513638Z","shell.execute_reply.started":"2026-01-09T18:58:59.509263Z","shell.execute_reply":"2026-01-09T18:58:59.512941Z"}},"outputs":[],"execution_count":null},{"id":"23f02d08","cell_type":"markdown","source":"# 15. Discussion: strengths, limitations, and extension paths\n\n## Strengths\n- Strong pretrained prior: RhoFold+ has been designed specifically for RNA and validated on diverse RNA benchmarks.\n- Correct handling of Kaggle constraints: no internet, fixed file paths, strict CSV format.\n- High quality engineering: caching, deterministic behavior, and explicit validation.\n- Robustness: fallback baseline ensures a valid submission even if external assets are unavailable.\n\n## Limitations\n- Single sequence mode may underperform on difficult targets where evolutionary coupling is essential.\n- Disabling refinement trades physical plausibility for speed.\n- The pipeline predicts only C1' coordinates; it does not explicitly model full atom geometry or steric clashes.\n\n## Extensions\n1. Enable MSA mode and use the provided MSAs directly. If the RhoFold parser expects A3M, improve the conversion by removing insertion columns relative to the query.\n2. Add lightweight refinement: apply a small number of gradient steps minimizing bond length and angle penalties, while preserving global topology.\n3. Add template discovery: build a larger template library from `PDB_RNA/` and run fast sequence searches to seed coordinates.\n4. Ensemble: combine multiple predictors and select the best among candidates using self consistency metrics.\n\nThese extensions can be implemented while keeping the notebook compliant, provided all external data is freely available and attached through Kaggle.\n","metadata":{}},{"id":"7df47ef2","cell_type":"code","source":"# Basic runtime report\nelapsed = time.time() - t0\nprint(f\"Total inference time: {elapsed:.1f} seconds\")\nprint(\"Submission path:\", out_path)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:59:03.802249Z","iopub.execute_input":"2026-01-09T18:59:03.803011Z","iopub.status.idle":"2026-01-09T18:59:03.806952Z","shell.execute_reply.started":"2026-01-09T18:59:03.802982Z","shell.execute_reply":"2026-01-09T18:59:03.806372Z"}},"outputs":[],"execution_count":null},{"id":"e0c53e26","cell_type":"markdown","source":"# 16. Appendix A: practical debugging checklist\n\nIf your submission scores unexpectedly low or fails, check:\n\n1. **File name**: must be exactly `submission.csv` in the working directory.\n2. **Row count**: must match `sample_submission.csv`.\n3. **IDs**: must match the template exactly. Do not reorder.\n4. **Residue indexing**: must be 1 based. Residue 1 is the first nucleotide.\n5. **Coordinate columns**: must include all of `x_1,y_1,z_1,...,x_5,y_5,z_5`.\n6. **Finite values**: no NaN, inf. Clip to [-999.999, 9999.999].\n7. **Runtime**: ensure the notebook completes within the configured limits.\n\nFor model debugging:\n- Start with `CFG.MAX_TARGETS = 2` to test end to end quickly.\n- Ensure RhoFold assets are attached and detected.\n- Inspect the cached output folder for each target: it should contain a PDB file.\n\nFor reproducibility:\n- Keep the seed fixed.\n- Avoid nondeterministic CuDNN behavior when sampling.\n","metadata":{}},{"id":"346e8474","cell_type":"code","source":"# List cache directory contents for quick debugging\nprint(\"CACHE_DIR:\", CACHE_DIR)\nfor p in sorted(CACHE_DIR.glob(\"*\"))[:5]:\n    print(\" -\", p.name)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:59:18.89204Z","iopub.execute_input":"2026-01-09T18:59:18.892481Z","iopub.status.idle":"2026-01-09T18:59:18.897275Z","shell.execute_reply.started":"2026-01-09T18:59:18.892451Z","shell.execute_reply":"2026-01-09T18:59:18.896638Z"}},"outputs":[],"execution_count":null},{"id":"4b773b8f","cell_type":"markdown","source":"# 17. Appendix B: full references\n\nThe following sources motivated design choices and provide the scientific context. For reproducibility, each entry includes enough bibliographic detail to locate the original work.\n\n1. **Stanford RNA 3D Folding Part 2 (Kaggle).** Competition overview, dataset schema, and submission requirements. (Kaggle competition page, accessed 2026-01-09).\n2. **US align.** S. Zhang et al. \"US align: Universal structure alignment of proteins, nucleic acids and macromolecular complexes.\" *Nature Methods* (2022). Describes the alignment procedure and TM score maximization used in evaluation.\n3. **TM score.** Y. Zhang and J. Skolnick. \"Scoring function for automated assessment of protein structure template quality.\" *Proteins* (2004). Introduces TM score and its scaling behavior.\n4. **RhoFold+.** T. Shen et al. \"Accurate RNA 3D structure prediction using a language model based deep learning approach.\" *Nature Methods* (2024). Presents RhoFold+ and its inference workflow.\n5. **Stanford RNA 3D Folding Part 1 outcome.** Preprint titled \"Template based RNA structure prediction advanced through a blind code competition.\" Summarizes competition insights and highlights the competitiveness of template oriented pipelines.\n6. **CASP16 assessors paper.** Assessment of RNA modeling performance in CASP16, providing benchmarking context for RNA structure prediction methods.\n\nBecause Kaggle disables internet during execution, these references are included only as documentation and are not accessed programmatically.\n","metadata":{}},{"id":"672ddb68","cell_type":"code","source":"# End of notebook marker\nprint(\"Notebook execution completed. submission.csv is ready.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-09T18:59:24.075529Z","iopub.execute_input":"2026-01-09T18:59:24.076145Z","iopub.status.idle":"2026-01-09T18:59:24.079641Z","shell.execute_reply.started":"2026-01-09T18:59:24.076116Z","shell.execute_reply":"2026-01-09T18:59:24.078883Z"}},"outputs":[],"execution_count":null}]}