{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","nbconvert_exporter":"python","pygments_lexer":"ipython3","version":"3.12.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"datasetVersion","sourceId":14604295,"datasetId":9328538,"databundleVersionId":15440074},{"sourceType":"datasetVersion","sourceId":14874339,"datasetId":9502242,"databundleVersionId":15736806},{"sourceType":"datasetVersion","sourceId":10855324,"datasetId":6742586,"databundleVersionId":11219268}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true},"papermill":{"default_parameters":{},"duration":165.02465,"end_time":"2026-03-02T08:06:01.924463","environment_variables":{},"exception":null,"input_path":"__notebook__.ipynb","output_path":"__notebook__.ipynb","parameters":{},"start_time":"2026-03-02T08:03:16.899813","version":"2.6.0"}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"f978dae1","cell_type":"markdown","source":"# Stanford RNA 3D Folding — Protenix + TBM Pipeline\n\n\n
What this notebook does\nThis notebook implements a
three-phase hybrid prediction pipeline for the Stanford RNA 3D Folding Kaggle competition. It combines classical Template-Based Modeling (TBM) with
Protenix — ByteDance Research's open-source reimplementation of AlphaFold3 — to predict the 3D atomic coordinates of RNA molecules.\n
\nThe key insight driving this design is that
TBM dramatically outperforms pure neural network prediction when close structural templates exist in the training data. As shown in the version table below, the combined approach scores
0.408 on the leaderboard versus
0.249 for Protenix alone.\n
\n
References and prior work:\n
\n
\n
\n
// AI · ML · Data Science
\n
Gourab Roy
\n
AI/ML practitioner shipping notebooks & code daily — built in the open.
\n
\n
\n
\nKey finding
\nTBM is not a fallback — it is the primary driver of performance. Pure Protenix without template initialization scores 0.159 lower, showing how much the neural network relies on structural priors to converge on a good fold.\n
\n\n\n\n\n| Version | \nDescription | \nLB Score | \n
\n\n\n\n| 3 | \nProtenix + TBM (current best) | \n0.408 | \n
\n\n| 4 | \nPure Protenix, no TBM | \n0.249 | \n
\n\n| 5 | \nUSE_MSA + USE_RNA_MSA | \nTBC | \n
\n\n| 6 | \nUSE_MSA only | \nTBC | \n
\n\n| 7 | \nUSE_RNA_MSA only | \nTBC | \n
\n\n
","metadata":{"papermill":{"duration":0.002917,"end_time":"2026-03-02T08:03:19.260043","exception":false,"start_time":"2026-03-02T08:03:19.257126","status":"completed"},"tags":[]}},{"id":"90626c3e","cell_type":"markdown","source":"## Part 1 — Environment Setup\n\n### Dependencies\n\n\nBiopython install
\nBiopython is required for pairwise sequence alignment. On Kaggle it is available as a pre-uploaded wheel to avoid internet access restrictions. The install line is commented out below since it may already be present in the environment. Uncomment if you get an import error for Bio.Align.\n
","metadata":{"papermill":{"duration":0.002932,"end_time":"2026-03-02T08:03:19.265829","exception":false,"start_time":"2026-03-02T08:03:19.262897","status":"completed"},"tags":[]}},{"id":"27525d0f","cell_type":"code","source":"# !pip install /kaggle/input/datasets/ogurtsov/biopython/biopython-1.85-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl","metadata":{"execution":{"iopub.status.busy":"2026-03-16T04:34:47.449723Z","iopub.execute_input":"2026-03-16T04:34:47.450514Z","iopub.status.idle":"2026-03-16T04:34:47.454576Z","shell.execute_reply.started":"2026-03-16T04:34:47.450482Z","shell.execute_reply":"2026-03-16T04:34:47.453952Z"},"papermill":{"duration":0.009173,"end_time":"2026-03-02T08:03:19.277804","exception":false,"start_time":"2026-03-02T08:03:19.268631","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"31030886","cell_type":"markdown","source":"### Competition vs. Local Mode\n\n\n
How the IS_KAGGLE flag works\n
IS_KAGGLE detects whether the notebook is running as an official Kaggle competition submission (when
KAGGLE_IS_COMPETITION_RERUN is set by Kaggle's infrastructure) or interactively during development.\n
\n - Local mode: only
LOCAL_N_SAMPLES sequences are processed, keeping iteration fast. \n - Competition mode: all test sequences are processed.
\n
\nThis single flag controls branching throughout the notebook — no manual commenting or uncommenting needed.\n
\n
Two environment variables are set before any imports\n
\n\n\n| Variable | \nValue | \nReason | \n
\n\n\n\nLAYERNORM_TYPE | \n\"torch\" | \nForces PyTorch's native layer normalisation instead of the custom CUDA kernel (fast_layernorm). Slower but works on P100 GPUs where the kernel may not be available. | \n
\n\nRNA_MSA_DEPTH_LIMIT | \n\"512\" | \nCaps the number of aligned sequences fed into the MSA module. Deeper MSAs improve accuracy but increase memory and runtime significantly. | \n
\n\n
\n
\n
What get_c1_mask does\nThis is an early definition of the
get_c1_mask utility, which is redefined with additional logic inside the main code block. It extracts a boolean mask identifying which atoms in Protenix's output correspond to
C1' atoms — the sugar carbons that serve as the representative coordinate for each nucleotide.\n
\nThe function tries three approaches in order of reliability:\n
\n - The
centre_atom_mask attribute on the atom array. \n - The
atom_name attribute. \n - A heuristic that picks whichever index value (11 or 12) in
atom_to_tokatom_idx produces a count closest to the sequence length. \n
\nThis defensive design handles API differences across Protenix versions.\n
\nThe entire pipeline is defined in this cell as a collection of functions culminating in main(). It is structured into five logical sections described below.\n
\n\n### 2.1 — Paths, Constants, and Configuration\n\n\nAll file paths and tunable parameters are defined at the top. Key values:\n
\n MAX_SEQ_LEN = 512 — sequences are truncated to this length before being passed to Protenix, preventing GPU out-of-memory errors. Full-length TBM predictions are still generated for all sequences regardless of length.MIN_SIMILARITY and MIN_PERCENT_IDENTITY — quality gates for template acceptance. Setting MIN_PERCENT_IDENTITY = 50.0 means only templates sharing at least half their nucleotide letters with the query are accepted.USE_PROTENIX = True — set to False to skip the neural network stage entirely and fall back to de novo structures, useful for rapid debugging.- MSA and template flags (
USE_MSA, USE_TEMPLATE, USE_RNA_MSA) are read from environment variables, enabling ablation experiments without code changes. \n
\n
\n
Sequence alignment uses a
PairwiseAligner with match score 2.0, mismatch penalty -1.5, and gap open -8. End gaps are penalised symmetrically on both query and target sides. The length-difference filter is 30%, so templates with very different lengths are rejected early.\n
\n
find_similar_sequences_detailed returns the top 30 candidates and computes two quality metrics per template: normalised alignment score and percent identity (fraction of aligned positions where query and template share the same nucleotide letter). Both are used as acceptance criteria.\n
\n
adapt_template_to_query maps template coordinates onto the query via the alignment's matched blocks, then fills gaps by linear interpolation. The fallback for edge residues with no valid anchor shifts by
[3, 0, 0] — a deterministic 3 Å step along the x-axis — making results fully reproducible.\n
\n
adaptive_rna_constraints applies three tiers of physical corrections per chain segment, run for multiple passes:\n
\n - Bond constraint (i to i+1): drives consecutive C1'-C1' distances toward 5.95 Å.
\n - Next-neighbour constraint (i to i+2): drives second-neighbour distances toward 10.2 Å, resisting sharp zigzag backbone conformations.
\n - Laplacian smoothing: nudges each internal residue toward the midpoint of its two neighbours, removing jagged kinks.
\n - Self-avoidance: for chains of 25+ residues, repels non-adjacent atoms closer than 3.2 Å. Uses a sampled subset of positions for chains longer than 220 residues to keep the O(n²) computation tractable.
\n
\nConstraint strength scales as
0.75 × (1 - min(confidence, 0.97)), meaning very high-similarity templates receive only minimal correction while de novo structures receive the full correction force.\n
\nRather than adding undifferentiated Gaussian noise to generate 5 diverse candidates, each prediction slot gets a qualitatively distinct structural perturbation:\n
\n - Slot 0: raw adapted template coordinates, unmodified.
\n - Slot 1: small Gaussian noise, magnitude proportional to
max(0.01, (0.40 - similarity) × 0.06). Better templates get less perturbation. \n - Slot 2: hinge rotation — a random pivot within the longest chain segment rotates all downstream residues by up to ±22°. This mimics a physical domain hinge motion.
\n - Slot 3: chain jitter — each chain segment is independently rotated by ±12° and translated by up to 1.5 Å, then the global centre of mass is restored. This simulates inter-chain conformational variability in multi-copy complexes.
\n - Slot 4: smooth wiggle — a smooth displacement field built from 6 interpolated random control points is applied to each chain, producing a global wave-like deformation.
\n
\n
\nProtenix only runs on sequences where TBM produced fewer than 5 predictions. It is asked to generate exactly as many samples as are still missing (N_sample = n_needed), rather than always generating 5 and discarding extras. This conserves GPU time significantly.\n
\nSequences are truncated to MAX_SEQ_LEN before being passed to Protenix. After inference, C1' coordinates are extracted from the raw all-atom output using the defensive masking logic described above. Extensive debug prints are retained in this section — added during development to diagnose coordinate extraction failures and detect model collapse (all residues converging to the same point).\n
\nIf Protenix fails for any sequence (exception, data error, or detected collapse), that target's slots remain empty and are filled in Phase 3.\n
\n\n### 2.5 — Combine and Save\n\n\nTBM predictions are taken first, Protenix predictions fill remaining slots in order, and any still-empty slots receive de novo helical structures. The final coordinates are clipped to [-999.999, 9999.999] to stay within the range the PDB format can represent — a requirement imposed by the competition's scoring metric.\n
","metadata":{"papermill":{"duration":0.002845,"end_time":"2026-03-02T08:03:23.536929","exception":false,"start_time":"2026-03-02T08:03:23.534084","status":"completed"},"tags":[]}},{"id":"1669cad3","cell_type":"code","source":"\n# ─────────────── Paths & Constants ───────────────────────────────────────────\nDATA_BASE = \"/kaggle/input/stanford-rna-3d-folding-2\"\nDEFAULT_TEST_CSV = f\"{DATA_BASE}/test_sequences.csv\"\nDEFAULT_TRAIN_CSV = f\"{DATA_BASE}/train_sequences.csv\"\nDEFAULT_TRAIN_LBLS = f\"{DATA_BASE}/train_labels.csv\"\nDEFAULT_VAL_CSV = f\"{DATA_BASE}/validation_sequences.csv\"\nDEFAULT_VAL_LBLS = f\"{DATA_BASE}/validation_labels.csv\"\nDEFAULT_OUTPUT = \"/kaggle/working/submission.csv\"\n\nDEFAULT_CODE_DIR = (\n \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted\"\n \"/Protenix-v1-adjust-v2/Protenix-v1-adjust-v2/Protenix-v1\"\n)\nDEFAULT_ROOT_DIR = DEFAULT_CODE_DIR\n\nMODEL_NAME = \"protenix_base_20250630_v1.0.0\"\nN_SAMPLE = 5\nSEED = 42\nMAX_SEQ_LEN = int(os.environ.get(\"MAX_SEQ_LEN\", \"512\"))\nCHUNK_OVERLAP = int(os.environ.get(\"CHUNK_OVERLAP\", \"64\"))\n\n# TBM quality thresholds — sequences below these get routed to Protenix\nMIN_SIMILARITY = float(os.environ.get(\"MIN_SIMILARITY\", \"0.0\"))\nMIN_PERCENT_IDENTITY = float(os.environ.get(\"MIN_PERCENT_IDENTITY\", \"50.0\"))\n\n# Set False to skip Protenix and use de-novo fallback instead\nUSE_PROTENIX = True\n\n\ndef parse_bool(value: str, default: bool = False) -> str:\n v = str(value).strip().lower()\n if v in {\"1\", \"true\", \"t\", \"yes\", \"y\", \"on\"}:\n return \"true\"\n if v in {\"0\", \"false\", \"f\", \"no\", \"n\", \"off\"}:\n return \"false\"\n return \"true\" if default else \"false\"\n\n\nUSE_MSA = parse_bool(os.environ.get(\"USE_MSA\", \"false\"))\nUSE_TEMPLATE = parse_bool(os.environ.get(\"USE_TEMPLATE\", \"false\"))\nUSE_RNA_MSA = parse_bool(os.environ.get(\"USE_RNA_MSA\", \"true\"))\n\nMODEL_N_SAMPLE = int(os.environ.get(\"MODEL_N_SAMPLE\", str(N_SAMPLE)))\n\n\n# ─────────────── General Utilities ───────────────────────────────────────────\ndef seed_everything(seed: int) -> None:\n os.environ[\"PYTHONHASHSEED\"] = str(seed)\n os.environ[\"CUBLAS_WORKSPACE_CONFIG\"] = \":4096:8\"\n torch.manual_seed(seed)\n torch.cuda.manual_seed(seed)\n torch.cuda.manual_seed_all(seed)\n np.random.seed(seed)\n torch.backends.cudnn.benchmark = False\n torch.backends.cudnn.deterministic = True\n torch.backends.cudnn.enabled = True\n torch.use_deterministic_algorithms(True)\n\n\ndef resolve_paths():\n test_csv = os.environ.get(\"TEST_CSV\", DEFAULT_TEST_CSV)\n output_csv = os.environ.get(\"SUBMISSION_CSV\", DEFAULT_OUTPUT)\n code_dir = os.environ.get(\"PROTENIX_CODE_DIR\", DEFAULT_CODE_DIR)\n root_dir = os.environ.get(\"PROTENIX_ROOT_DIR\", DEFAULT_ROOT_DIR)\n return test_csv, output_csv, code_dir, root_dir\n\n\ndef ensure_required_files(root_dir: str) -> None:\n for p, name in [\n (Path(root_dir) / \"checkpoint\" / f\"{MODEL_NAME}.pt\", \"checkpoint\"),\n (Path(root_dir) / \"common\" / \"components.cif\", \"CCD file\"),\n (Path(root_dir) / \"common\" / \"components.cif.rdkit_mol.pkl\", \"CCD cache\"),\n ]:\n if not p.exists():\n raise FileNotFoundError(f\"Missing {name}: {p}\")\n\n\n# ─────────────── Protenix Input / Config Helpers ─────────────────────────────\ndef build_input_json(df: pd.DataFrame, json_path: str) -> None:\n data = [\n {\n \"name\": row[\"target_id\"],\n \"covalent_bonds\": [],\n \"sequences\": [{\"rnaSequence\": {\"sequence\": row[\"sequence\"], \"count\": 1}}],\n }\n for _, row in df.iterrows()\n ]\n with open(json_path, \"w\", encoding=\"utf-8\") as f:\n json.dump(data, f)\n\n\ndef build_configs(input_json_path: str, dump_dir: str, model_name: str):\n from configs.configs_base import configs as configs_base\n from configs.configs_data import data_configs\n from configs.configs_inference import inference_configs\n from configs.configs_model_type import model_configs\n from protenix.config.config import parse_configs\n\n base = {**configs_base, **{\"data\": data_configs}, **inference_configs}\n\n def deep_update(t, p):\n for k, v in p.items():\n if isinstance(v, dict) and k in t and isinstance(t[k], dict):\n deep_update(t[k], v)\n else:\n t[k] = v\n\n deep_update(base, model_configs[model_name])\n arg_str = \" \".join([\n f\"--model_name {model_name}\",\n f\"--input_json_path {input_json_path}\",\n f\"--dump_dir {dump_dir}\",\n f\"--use_msa {USE_MSA}\",\n f\"--use_template {USE_TEMPLATE}\",\n f\"--use_rna_msa {USE_RNA_MSA}\",\n f\"--sample_diffusion.N_sample {MODEL_N_SAMPLE}\",\n f\"--seeds {SEED}\",\n ])\n return parse_configs(configs=base, arg_str=arg_str, fill_required_with_null=True)\n\n\ndef get_c1_mask(data: dict, atom_array) -> torch.Tensor:\n # 1. Try atom_array attributes first\n if atom_array is not None:\n try:\n if hasattr(atom_array, \"centre_atom_mask\"):\n m = atom_array.centre_atom_mask == 1\n if hasattr(atom_array, \"is_rna\"):\n m = m & atom_array.is_rna\n return torch.from_numpy(m).bool()\n \n if hasattr(atom_array, \"atom_name\"):\n base = atom_array.atom_name == \"C1'\"\n if hasattr(atom_array, \"is_rna\"):\n base = base & atom_array.is_rna\n return torch.from_numpy(base).bool()\n except Exception:\n pass\n\n # 2. Fallback to feature dict\n f = data[\"input_feature_dict\"]\n \n # CASE A: center_atom_mask exists\n if \"center_atom_mask\" in f:\n return (f[\"center_atom_mask\"] == 1).bool()\n if \"centre_atom_mask\" in f:\n return (f[\"centre_atom_mask\"] == 1).bool()\n \n # CASE B: Use atom_name\n if \"atom_name\" in f:\n # Check against \"C1'\" (byte encoded or string?)\n # For now assume typical behavior is center_atom_mask is present.\n pass\n\n # CASE C: atom_to_tokatom_idx fallback\n # The index for C1' is typically 11 or 12 depending on featurizer.\n # Let's try to match exactly C1' if possible.\n # But usually 'centre_atom_mask' should be there.\n \n # If we fall through, assume standard mask\n return (f[\"atom_to_tokatom_idx\"] == 11).bool()\n\n\ndef get_feature_c1_mask(data: dict) -> torch.Tensor:\n f = data[\"input_feature_dict\"]\n if \"centre_atom_mask\" in f:\n return f[\"centre_atom_mask\"].long() == 1\n return f[\"atom_to_tokatom_idx\"].long() == 12\n\n\ndef coords_to_rows(target_id: str, seq: str, coords: np.ndarray) -> list:\n \"\"\"coords shape: (N_SAMPLE, seq_len, 3)\"\"\"\n rows = []\n for i in range(len(seq)):\n row = {\"ID\": f\"{target_id}_{i + 1}\", \"resname\": seq[i], \"resid\": i + 1}\n for s in range(N_SAMPLE):\n if s < coords.shape[0] and i < coords.shape[1]:\n x, y, z = coords[s, i]\n else:\n x, y, z = 0.0, 0.0, 0.0\n row[f\"x_{s + 1}\"] = float(x)\n row[f\"y_{s + 1}\"] = float(y)\n row[f\"z_{s + 1}\"] = float(z)\n rows.append(row)\n return rows\n\n\ndef pad_samples(coords: np.ndarray, n: int) -> np.ndarray:\n if coords.shape[0] >= n:\n return coords[:n]\n if coords.shape[0] == 0:\n return np.zeros((n, coords.shape[1], 3), dtype=coords.dtype)\n extra = np.repeat(coords[:1], n - coords.shape[0], axis=0)\n return np.concatenate([coords, extra], axis=0)\n\n\n# ─────────────── TBM Core Functions ──────────────────────────────────────────\ndef _make_aligner() -> PairwiseAligner:\n al = PairwiseAligner()\n al.mode = \"global\"\n al.match_score = 2\n al.mismatch_score = -1.5\n al.open_gap_score = -8\n al.extend_gap_score = -0.4\n al.query_left_open_gap_score = -8\n al.query_left_extend_gap_score = -0.4\n al.query_right_open_gap_score = -8\n al.query_right_extend_gap_score = -0.4\n al.target_left_open_gap_score = -8\n al.target_left_extend_gap_score = -0.4\n al.target_right_open_gap_score = -8\n al.target_right_extend_gap_score = -0.4\n return al\n\n\n_aligner = _make_aligner()\n\n\ndef parse_stoichiometry(stoich: str) -> list:\n if pd.isna(stoich) or str(stoich).strip() == \"\":\n return []\n return [(ch.strip(), int(cnt)) for part in str(stoich).split(\";\")\n for ch, cnt in [part.split(\":\")]]\n\n\ndef parse_fasta(fasta_content: str) -> dict:\n out, cur, parts = {}, None, []\n for line in str(fasta_content).splitlines():\n line = line.strip()\n if not line:\n continue\n if line.startswith(\">\"):\n if cur is not None:\n out[cur] = \"\".join(parts)\n cur = line[1:].split()[0]\n parts = []\n else:\n parts.append(line.replace(\" \", \"\"))\n if cur is not None:\n out[cur] = \"\".join(parts)\n return out\n\n\ndef get_chain_segments(row) -> list:\n seq = row[\"sequence\"]\n stoich = row.get(\"stoichiometry\", \"\")\n all_sq = row.get(\"all_sequences\", \"\")\n if (pd.isna(stoich) or pd.isna(all_sq)\n or str(stoich).strip() == \"\" or str(all_sq).strip() == \"\"):\n return [(0, len(seq))]\n try:\n chain_dict = parse_fasta(all_sq)\n order = parse_stoichiometry(stoich)\n segs, pos = [], 0\n for ch, cnt in order:\n base = chain_dict.get(ch)\n if base is None:\n return [(0, len(seq))]\n for _ in range(cnt):\n segs.append((pos, pos + len(base)))\n pos += len(base)\n return segs if pos == len(seq) else [(0, len(seq))]\n except Exception:\n return [(0, len(seq))]\n\n\ndef build_segments_map(df: pd.DataFrame) -> tuple:\n seg_map, stoich_map = {}, {}\n for _, r in df.iterrows():\n tid = r[\"target_id\"]\n seg_map[tid] = get_chain_segments(r)\n raw_s = r.get(\"stoichiometry\", \"\")\n stoich_map[tid] = \"\" if pd.isna(raw_s) else str(raw_s)\n return seg_map, stoich_map\n\n\ndef process_labels(labels_df: pd.DataFrame) -> dict:\n coords = {}\n prefixes = labels_df[\"ID\"].str.rsplit(\"_\", n=1).str[0]\n for prefix, grp in labels_df.groupby(prefixes):\n coords[prefix] = grp.sort_values(\"resid\")[[\"x_1\", \"y_1\", \"z_1\"]].values\n return coords\n\n\ndef _build_aligned_strings(query_seq, template_seq, alignment):\n q_segs, t_segs = alignment.aligned\n aq, at, qi, ti = [], [], 0, 0\n for (qs, qe), (ts, te) in zip(q_segs, t_segs):\n while qi < qs: aq.append(query_seq[qi]); at.append(\"-\"); qi += 1\n while ti < ts: aq.append(\"-\"); at.append(template_seq[ti]); ti += 1\n for qp, tp in zip(range(qs, qe), range(ts, te)):\n aq.append(query_seq[qp]); at.append(template_seq[tp])\n qi, ti = qe, te\n while qi < len(query_seq): aq.append(query_seq[qi]); at.append(\"-\"); qi += 1\n while ti < len(template_seq): aq.append(\"-\"); at.append(template_seq[ti]); ti += 1\n return \"\".join(aq), \"\".join(at)\n\n\ndef find_similar_sequences_detailed(query_seq, train_seqs_df, train_coords_dict, top_n=30):\n results = []\n for _, row in train_seqs_df.iterrows():\n tid, tseq = row[\"target_id\"], row[\"sequence\"]\n if tid not in train_coords_dict:\n continue\n if abs(len(tseq) - len(query_seq)) / max(len(tseq), len(query_seq)) > 0.3:\n continue\n aln = next(iter(_aligner.align(query_seq, tseq)))\n norm_s = aln.score / (2 * min(len(query_seq), len(tseq)))\n identical = sum(\n 1 for (qs, qe), (ts, te) in zip(*aln.aligned)\n for qp, tp in zip(range(qs, qe), range(ts, te))\n if query_seq[qp] == tseq[tp]\n )\n pct_id = 100 * identical / len(query_seq)\n aq, at = _build_aligned_strings(query_seq, tseq, aln)\n results.append((tid, tseq, norm_s, train_coords_dict[tid], pct_id, aq, at))\n results.sort(key=lambda x: x[2], reverse=True)\n return results[:top_n]\n\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords) -> np.ndarray:\n aln = next(iter(_aligner.align(query_seq, template_seq)))\n new_coords = np.full((len(query_seq), 3), np.nan)\n for (qs, qe), (ts, te) in zip(*aln.aligned):\n chunk = template_coords[ts:te]\n if len(chunk) == (qe - qs):\n new_coords[qs:qe] = chunk\n for i in range(len(new_coords)):\n if np.isnan(new_coords[i, 0]):\n pv = next((j for j in range(i - 1, -1, -1) if not np.isnan(new_coords[j, 0])), -1)\n nv = next((j for j in range(i + 1, len(new_coords)) if not np.isnan(new_coords[j, 0])), -1)\n if pv >= 0 and nv >= 0:\n w = (i - pv) / (nv - pv)\n new_coords[i] = (1 - w) * new_coords[pv] + w * new_coords[nv]\n elif pv >= 0:\n new_coords[i] = new_coords[pv] + [3, 0, 0]\n elif nv >= 0:\n new_coords[i] = new_coords[nv] + [3, 0, 0]\n else:\n new_coords[i] = [i * 3, 0, 0]\n return np.nan_to_num(new_coords)\n\n\ndef adaptive_rna_constraints(coords, target_id, segments_map, confidence=1.0, passes=2) -> np.ndarray:\n X = coords.copy()\n segments = segments_map.get(target_id, [(0, len(X))])\n strength = max(0.75 * (1.0 - min(confidence, 0.97)), 0.02)\n for _ in range(passes):\n for s, e in segments:\n C = X[s:e]; L = e - s\n if L < 3:\n continue\n # bond i–i+1 ~5.95 Å\n d = C[1:] - C[:-1]; dist = np.linalg.norm(d, axis=1) + 1e-6\n adj = d * ((5.95 - dist) / dist)[:, None] * (0.22 * strength)\n C[:-1] -= adj; C[1:] += adj\n # soft i–i+2 ~10.2 Å\n d2 = C[2:] - C[:-2]; d2n = np.linalg.norm(d2, axis=1) + 1e-6\n adj2 = d2 * ((10.2 - d2n) / d2n)[:, None] * (0.10 * strength)\n C[:-2] -= adj2; C[2:] += adj2\n # Laplacian smoothing\n C[1:-1] += (0.06 * strength) * (0.5 * (C[:-2] + C[2:]) - C[1:-1])\n # self-avoidance\n if L >= 25:\n idx = np.linspace(0, L - 1, min(L, 160)).astype(int) if L > 220 else np.arange(L)\n P = C[idx]; diff = P[:, None, :] - P[None, :, :]\n dm = np.linalg.norm(diff, axis=2) + 1e-6\n sep = np.abs(idx[:, None] - idx[None, :])\n mask = (sep > 2) & (dm < 3.2)\n if np.any(mask):\n vec = (diff * ((3.2 - dm) / dm)[:, :, None] * mask[:, :, None]).sum(axis=1)\n C[idx] += (0.015 * strength) * vec\n X[s:e] = C\n return X\n\n\ndef _rotmat(axis, ang):\n a = np.asarray(axis, float); a /= np.linalg.norm(a) + 1e-12\n x, y, z = a; c, s = np.cos(ang), np.sin(ang); CC = 1 - c\n return np.array([[c+x*x*CC, x*y*CC-z*s, x*z*CC+y*s],\n [y*x*CC+z*s, c+y*y*CC, y*z*CC-x*s],\n [z*x*CC-y*s, z*y*CC+x*s, c+z*z*CC]])\n\n\ndef apply_hinge(coords, seg, rng, deg=22):\n s, e = seg; L = e - s\n if L < 30: return coords\n pivot = s + int(rng.integers(10, L - 10))\n R = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-deg, deg))))\n X = coords.copy(); p0 = X[pivot].copy()\n X[pivot+1:e] = (X[pivot+1:e] - p0) @ R.T + p0\n return X\n\n\ndef jitter_chains(coords, segs, rng, deg=12, trans=1.5):\n X = coords.copy(); gc_ = X.mean(0, keepdims=True)\n for s, e in segs:\n R = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-deg, deg))))\n shift = rng.normal(size=3); shift = shift / (np.linalg.norm(shift) + 1e-12) * float(rng.uniform(0, trans))\n c = X[s:e].mean(0, keepdims=True)\n X[s:e] = (X[s:e] - c) @ R.T + c + shift\n X -= X.mean(0, keepdims=True) - gc_\n return X\n\n\ndef smooth_wiggle(coords, segs, rng, amp=0.8):\n X = coords.copy()\n for s, e in segs:\n L = e - s\n if L < 20: continue\n ctrl = np.linspace(0, L - 1, 6); disp = rng.normal(0, amp, (6, 3)); t = np.arange(L)\n X[s:e] += np.vstack([np.interp(t, ctrl, disp[:, k]) for k in range(3)]).T\n return X\n\n\ndef generate_rna_structure(sequence: str, seed=None) -> np.ndarray:\n \"\"\"Idealized A-form RNA helix — last-resort de-novo fallback.\"\"\"\n if seed is not None:\n np.random.seed(seed)\n n = len(sequence); coords = np.zeros((n, 3))\n for i in range(n):\n ang = i * 0.6\n coords[i] = [10.0 * np.cos(ang), 10.0 * np.sin(ang), i * 2.5]\n return coords\n\n\n# ─────────────── TBM Phase ───────────────────────────────────────────────────\ndef tbm_phase(test_df, train_seqs_df, train_coords_dict, segments_map):\n \"\"\"\n Phase 1 — Template-Based Modeling.\n\n Returns\n -------\n template_predictions : {target_id: [np.ndarray(seq_len, 3), ...]}\n 0 to N_SAMPLE predictions per target, from real templates.\n protenix_queue : {target_id: (n_needed, full_sequence)}\n Targets that still need more predictions.\n \"\"\"\n print(f\"\\n{'='*60}\")\n print(f\"PHASE 1: Template-Based Modeling\")\n print(f\" MIN_SIMILARITY = {MIN_SIMILARITY} | MIN_PCT_IDENTITY = {MIN_PERCENT_IDENTITY}\")\n print(f\"{'='*60}\")\n t0 = time.time()\n\n template_predictions: dict = {}\n protenix_queue: dict = {}\n\n for _, row in test_df.iterrows():\n tid = row[\"target_id\"]\n seq = row[\"sequence\"]\n segs = segments_map.get(tid, [(0, len(seq))])\n\n similar = find_similar_sequences_detailed(seq, train_seqs_df, train_coords_dict, top_n=30)\n preds = []\n used = set()\n\n for i, (tmpl_id, tmpl_seq, sim, tmpl_coords, pct_id, _, _) in enumerate(similar):\n if len(preds) >= N_SAMPLE:\n break\n if sim < MIN_SIMILARITY or pct_id < MIN_PERCENT_IDENTITY:\n break # list is sorted by sim, so no point continuing\n if tmpl_id in used:\n continue\n\n rng = np.random.default_rng((abs(hash(tid)) + i * 10007) % (2**32))\n adapted = adapt_template_to_query(seq, tmpl_seq, tmpl_coords)\n\n # Diversity transforms (same strategy as the 0-409 TBM notebook)\n slot = len(preds)\n if slot == 0:\n X = adapted\n elif slot == 1:\n X = adapted + rng.normal(0, max(0.01, (0.40 - sim) * 0.06), adapted.shape)\n elif slot == 2:\n longest = max(segs, key=lambda se: se[1] - se[0])\n X = apply_hinge(adapted, longest, rng)\n elif slot == 3:\n X = jitter_chains(adapted, segs, rng)\n else:\n X = smooth_wiggle(adapted, segs, rng)\n\n refined = adaptive_rna_constraints(X, tid, segments_map, confidence=sim)\n preds.append(refined)\n used.add(tmpl_id)\n\n template_predictions[tid] = preds\n n_needed = N_SAMPLE - len(preds)\n if n_needed > 0:\n protenix_queue[tid] = (n_needed, seq)\n print(f\" {tid} ({len(seq)} nt): {len(preds)} TBM → need {n_needed} from Protenix\")\n else:\n print(f\" {tid} ({len(seq)} nt): all {N_SAMPLE} from TBM ✓\")\n\n elapsed = time.time() - t0\n n_full = len(test_df) - len(protenix_queue)\n print(f\"\\nPhase 1 done in {elapsed:.1f}s\")\n print(f\" Fully covered by TBM : {n_full}\")\n print(f\" Need Protenix : {len(protenix_queue)}\")\n return template_predictions, protenix_queue\n\n\n# ─────────────── Main ────────────────────────────────────────────────────────\ndef main() -> None:\n test_csv, output_csv, code_dir, root_dir = resolve_paths()\n\n if not os.path.isdir(code_dir):\n raise FileNotFoundError(\n f\"Missing PROTENIX_CODE_DIR: {code_dir}. \"\n \"Set PROTENIX_CODE_DIR to the repo path.\"\n )\n\n os.environ[\"PROTENIX_ROOT_DIR\"] = root_dir\n sys.path.append(code_dir)\n ensure_required_files(root_dir)\n seed_everything(SEED)\n\n # ── Load test data ──────────────────────────────────────────────────────\n test_df_full = pd.read_csv(test_csv)\n test_df = (test_df_full.head(LOCAL_N_SAMPLES) if not IS_KAGGLE\n else test_df_full).reset_index(drop=True)\n print(f\"Test targets : {len(test_df)}\"\n + (\" (LOCAL MODE)\" if not IS_KAGGLE else \"\"))\n\n seq_by_id = dict(zip(test_df[\"target_id\"], test_df[\"sequence\"]))\n\n # Truncated copy for Protenix (Protenix has token limits)\n test_df_trunc = test_df.copy()\n test_df_trunc[\"sequence\"] = test_df_trunc[\"sequence\"].str[:MAX_SEQ_LEN]\n\n # ── Load training data for TBM ──────────────────────────────────────────\n print(\"\\nLoading training data for TBM …\")\n train_seqs = pd.read_csv(DEFAULT_TRAIN_CSV)\n val_seqs = pd.read_csv(DEFAULT_VAL_CSV)\n train_labels = pd.read_csv(DEFAULT_TRAIN_LBLS)\n val_labels = pd.read_csv(DEFAULT_VAL_LBLS)\n\n combined_seqs = pd.concat([train_seqs, val_seqs], ignore_index=True)\n combined_labels = pd.concat([train_labels, val_labels], ignore_index=True)\n train_coords = process_labels(combined_labels)\n segments_map, _ = build_segments_map(test_df)\n\n print(f\"Template pool: {len(combined_seqs)} sequences, {len(train_coords)} structures\")\n\n # ─── PHASE 1: TBM ──────────────────────────────────────────────────────\n template_preds, protenix_queue = tbm_phase(\n test_df, combined_seqs, train_coords, segments_map\n )\n\n # ─── PHASE 2: Protenix (only for targets that need extra predictions) ──\n protenix_preds: dict = {} # target_id -> np.ndarray (n_needed, seq_len, 3)\n\n if protenix_queue and USE_PROTENIX:\n print(f\"\\n{'='*60}\")\n print(f\"PHASE 2: Protenix for {len(protenix_queue)} targets\")\n print(f\"{'='*60}\")\n\n work_dir = Path(\"/kaggle/working\")\n work_dir.mkdir(parents=True, exist_ok=True)\n\n # Build input JSON only for queued targets\n queue_df = (test_df_trunc[test_df_trunc[\"target_id\"].isin(protenix_queue)]\n .reset_index(drop=True))\n input_json_path = str(work_dir / \"protenix_queue_input.json\")\n build_input_json(queue_df, input_json_path)\n\n from protenix.data.inference.infer_dataloader import InferenceDataset\n from runner.inference import (InferenceRunner,\n update_gpu_compatible_configs,\n update_inference_configs)\n\n configs = build_configs(input_json_path, str(work_dir / \"outputs\"), MODEL_NAME)\n configs = update_gpu_compatible_configs(configs)\n runner = InferenceRunner(configs)\n dataset = InferenceDataset(configs)\n\n for i in tqdm(range(len(dataset)), desc=\"Protenix\"):\n data, atom_array, error_message = dataset[i]\n target_id = data.get(\"sample_name\", f\"sample_{i}\")\n\n if target_id not in protenix_queue:\n continue\n\n n_needed, full_seq = protenix_queue[target_id]\n\n if error_message:\n print(f\" {target_id}: data error — {error_message}\")\n protenix_preds[target_id] = None\n del data, atom_array, error_message\n gc.collect(); torch.cuda.empty_cache(); gc.collect()\n continue\n\n try:\n new_cfg = update_inference_configs(configs, data[\"N_token\"].item())\n # Only generate as many samples as we actually need to fill the slots\n new_cfg.sample_diffusion.N_sample = n_needed\n runner.update_model_configs(new_cfg)\n\n prediction = runner.predict(data)\n raw_coords = prediction[\"coordinate\"] # Shape: [N_sample, all_atoms, 3]\n\n # -----------------------------------------------------------\n # DEBUG PRINT START\n # -----------------------------------------------------------\n print(f\"\\n[DEBUG] {target_id} | n_needed: {n_needed} | SeqLen: {len(full_seq)}\")\n print(f\"[DEBUG] raw_coords shape: {raw_coords.shape}\")\n \n feat = data[\"input_feature_dict\"]\n \n # Check potential masks\n mask_candidates = {}\n if \"centre_atom_mask\" in feat:\n m = feat[\"centre_atom_mask\"]\n mask_candidates['centre_atom_mask'] = (m.sum().item(), m.shape)\n \n if \"atom_to_tokatom_idx\" in feat:\n idx_11 = (feat[\"atom_to_tokatom_idx\"] == 11).sum().item()\n idx_12 = (feat[\"atom_to_tokatom_idx\"] == 12).sum().item()\n mask_candidates['idx_11'] = idx_11\n mask_candidates['idx_12'] = idx_12\n \n print(f\"[DEBUG] Mask candidates counts: {mask_candidates}\")\n # -----------------------------------------------------------\n # DEBUG PRINT END\n # -----------------------------------------------------------\n\n # ─────────────────────────────────────────────────────────────\n # DEBUG / FIX: Explicit C1' masking logic\n # ─────────────────────────────────────────────────────────────\n # Try to use 'centre_atom_mask' from features if possible\n if \"centre_atom_mask\" in feat:\n mask = (feat[\"centre_atom_mask\"] == 1).to(raw_coords.device)\n elif \"atom_to_tokatom_idx\" in feat:\n # Heuristic: pick the one closest to sequence length\n m11 = (feat[\"atom_to_tokatom_idx\"] == 11).to(raw_coords.device)\n m12 = (feat[\"atom_to_tokatom_idx\"] == 12).to(raw_coords.device)\n \n c11, c12 = m11.sum(), m12.sum()\n target_len = len(full_seq) # closer to N_token usually\n \n if abs(c11 - target_len) < abs(c12 - target_len):\n mask = m11\n print(f\"[DEBUG] Selected idx 11 mask (count={c11})\")\n else:\n mask = m12\n print(f\"[DEBUG] Selected idx 12 mask (count={c12})\")\n else:\n # Should not happen\n mask = torch.zeros(raw_coords.shape[1], dtype=torch.bool, device=raw_coords.device)\n \n # Extract\n coords = raw_coords[:, mask, :].detach().cpu().numpy()\n print(f\"[DEBUG] Extracted coords shape: {coords.shape}\")\n\n # If we get duplicate coordinates (collapsed), this is bad.\n # Check for duplications in first sample\n if coords.shape[1] > 1:\n diffs = np.linalg.norm(coords[0, 1:] - coords[0, :-1], axis=-1)\n if np.all(diffs < 1e-4):\n print(f\" WARNING: {target_id} has identical coordinates for all residues! (Model collapse?)\")\n \n # Pad/trim to full (un-truncated) sequence length\n if coords.shape[1] != len(full_seq):\n # Check for broadcast issue or model collapse\n if coords.shape[1] == 1 and len(full_seq) > 1:\n # Model outputted only 1 residue/atom but we need many?\n # Broadcast the single coord to all positions just in case (though highly suspicious)\n # Or perhaps mask was wrong and selected only 1 atom.\n # Do NOT broadcast, fill with zeros to be safe.\n print(f\"[DEBUG] WARNING: {target_id}: mask selected only 1 atom, but sequence is {len(full_seq)}\")\n # padded = np.zeros(...) -> kept as zeros\n else:\n padded = np.zeros((coords.shape[0], len(full_seq), 3), dtype=np.float32)\n min_len = min(coords.shape[1], len(full_seq))\n if min_len > 0:\n padded[:, :min_len, :] = coords[:, :min_len, :]\n coords = padded\n\n # Final check for identical coordinates (indicative of model failure)\n if coords.shape[1] > 1:\n diffs = np.linalg.norm(coords[0, 1:] - coords[0, :-1], axis=-1)\n if np.all(diffs < 1e-4):\n print(f\" WARNING: {target_id}: Identical coordinates detected! Resetting to zeros.\")\n coords = np.zeros_like(coords)\n\n protenix_preds[target_id] = coords\n print(f\" {target_id}: {coords.shape[0]} Protenix predictions generated\")\n\n except Exception as exc:\n print(f\" {target_id}: Protenix FAILED — {exc}\")\n import traceback\n traceback.print_exc()\n protenix_preds[target_id] = None\n\n finally:\n del prediction, raw_coords, mask, data, atom_array\n gc.collect(); torch.cuda.empty_cache(); gc.collect()\n# ...existing code...\n\n elif protenix_queue and not USE_PROTENIX:\n print(f\"\\nPHASE 2 skipped (USE_PROTENIX=False). \"\n f\"De-novo fallback will cover {len(protenix_queue)} targets.\")\n\n # ─── PHASE 3: Combine everything ───────────────────────────────────────\n print(f\"\\n{'='*60}\")\n print(\"PHASE 3: Combine TBM + Protenix + de-novo fallback\")\n print(f\"{'='*60}\")\n\n all_rows = []\n\n for _, row in test_df.iterrows():\n tid = row[\"target_id\"]\n seq = row[\"sequence\"]\n\n combined: list = list(template_preds.get(tid, [])) # TBM predictions\n\n # Append Protenix predictions to fill remaining slots\n ptx = protenix_preds.get(tid)\n if ptx is not None and ptx.ndim == 3:\n for j in range(ptx.shape[0]):\n if len(combined) >= N_SAMPLE:\n break\n combined.append(ptx[j]) # (seq_len, 3)\n\n # De-novo fallback for any still-empty slots\n n_denovo = 0\n while len(combined) < N_SAMPLE:\n seed_val = hash(tid) % 10000 + len(combined) * 1000\n dn = generate_rna_structure(seq, seed=seed_val)\n combined.append(adaptive_rna_constraints(dn, tid, segments_map, confidence=0.2))\n n_denovo += 1\n\n if n_denovo:\n print(f\" {tid}: {n_denovo} slot(s) filled with de-novo fallback\")\n\n # Stack to (N_SAMPLE, seq_len, 3) and write rows\n stacked = np.stack(combined[:N_SAMPLE], axis=0)\n all_rows.extend(coords_to_rows(tid, seq, stacked))\n\n # ── Save ───────────────────────────────────────────────────────────────\n sub = pd.DataFrame(all_rows)\n cols = [\"ID\", \"resname\", \"resid\"] + [\n f\"{c}_{i}\" for i in range(1, N_SAMPLE + 1) for c in [\"x\", \"y\", \"z\"]\n ]\n coord_cols = [c for c in cols if c.startswith((\"x_\", \"y_\", \"z_\"))]\n sub[coord_cols] = sub[coord_cols].clip(-999.999, 9999.999)\n sub[cols].to_csv(output_csv, index=False)\n\n print(f\"\\n✓ Saved submission to {output_csv} ({len(sub):,} rows)\")\n\n","metadata":{"execution":{"iopub.status.busy":"2026-03-16T04:34:51.248137Z","iopub.status.idle":"2026-03-16T04:34:51.248375Z","shell.execute_reply.started":"2026-03-16T04:34:51.248267Z","shell.execute_reply":"2026-03-16T04:34:51.248281Z"},"papermill":{"duration":0.08427,"end_time":"2026-03-02T08:03:23.624011","exception":false,"start_time":"2026-03-02T08:03:23.539741","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"a4e1eb12","cell_type":"markdown","source":"## Part 3 — Entry Point\n\n\nThe main() function defined above is called here. The if __name__ == \"__main__\" guard is preserved for compatibility with running the file as a standard Python script. When run as a Kaggle notebook cell, the guard evaluates to True and main() executes immediately.\n
","metadata":{"papermill":{"duration":0.002945,"end_time":"2026-03-02T08:03:23.630081","exception":false,"start_time":"2026-03-02T08:03:23.627136","status":"completed"},"tags":[]}},{"id":"9fa592b8","cell_type":"code","source":"\nif __name__ == \"__main__\":\n main()\n","metadata":{"execution":{"iopub.status.busy":"2026-03-16T04:34:51.249692Z","iopub.status.idle":"2026-03-16T04:34:51.250007Z","shell.execute_reply.started":"2026-03-16T04:34:51.249836Z","shell.execute_reply":"2026-03-16T04:34:51.249856Z"},"papermill":{"duration":154.940035,"end_time":"2026-03-02T08:05:58.573765","exception":false,"start_time":"2026-03-02T08:03:23.63373","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null},{"id":"8160ed2f","cell_type":"markdown","source":"## Part 4 — Verify Output\n\n\nSanity check
\nA quick check to confirm the submission file was written correctly and has the expected structure. Each row represents one nucleotide of one RNA sequence, with columns for the 5 predicted (x, y, z) coordinate sets.\n
","metadata":{"papermill":{"duration":0.003926,"end_time":"2026-03-02T08:05:58.581703","exception":false,"start_time":"2026-03-02T08:05:58.577777","status":"completed"},"tags":[]}},{"id":"d03d491a","cell_type":"code","source":"#read submission.csv\nsubmission_path = \"/kaggle/working/submission.csv\"\nsubmission_df = pd.read_csv(submission_path)\nprint(submission_df.head(20))","metadata":{"execution":{"iopub.status.busy":"2026-03-16T04:34:51.251142Z","iopub.status.idle":"2026-03-16T04:34:51.25144Z","shell.execute_reply.started":"2026-03-16T04:34:51.251274Z","shell.execute_reply":"2026-03-16T04:34:51.251297Z"},"papermill":{"duration":0.023399,"end_time":"2026-03-02T08:05:58.608796","exception":false,"start_time":"2026-03-02T08:05:58.585397","status":"completed"},"tags":[],"trusted":true},"outputs":[],"execution_count":null}]}