{"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,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":14874339,"datasetId":9502242,"databundleVersionId":15736806},{"sourceType":"datasetVersion","sourceId":10855324,"datasetId":6742586,"databundleVersionId":11219268}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os\nimport sys\nimport pandas as pd\n\nIS_KAGGLE = bool(os.environ.get(\"KAGGLE_IS_COMPETITION_RERUN\", \"\"))\nLOCAL_N_SAMPLES = 2\n\nif IS_KAGGLE:\n    print(\"Running in KAGGLE COMPETITION mode — all test targets will be processed.\")\nelse:\n    print(f\"Running in LOCAL mode — only the first {LOCAL_N_SAMPLES} test targets will be processed.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:55.925672Z","iopub.execute_input":"2026-02-20T12:02:55.926567Z","iopub.status.idle":"2026-02-20T12:02:55.931089Z","shell.execute_reply.started":"2026-02-20T12:02:55.926532Z","shell.execute_reply":"2026-02-20T12:02:55.930416Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# !pip install biopython","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:55.937214Z","iopub.execute_input":"2026-02-20T12:02:55.937468Z","iopub.status.idle":"2026-02-20T12:02:55.944982Z","shell.execute_reply.started":"2026-02-20T12:02:55.937448Z","shell.execute_reply":"2026-02-20T12:02:55.94429Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\nimport json\nimport os\nimport time\n\nos.environ[\"LAYERNORM_TYPE\"] = \"torch\"\nos.environ.setdefault(\"RNA_MSA_DEPTH_LIMIT\", \"512\")\n\nimport sys\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom Bio.Align import PairwiseAligner\nfrom tqdm import tqdm\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:55.946038Z","iopub.execute_input":"2026-02-20T12:02:55.946249Z","iopub.status.idle":"2026-02-20T12:02:55.957657Z","shell.execute_reply.started":"2026-02-20T12:02:55.946229Z","shell.execute_reply":"2026-02-20T12:02:55.957022Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_c1_mask(data: dict, atom_array=None) -> torch.Tensor:\n    \"\"\"Robustly extract the C1' atom mask from Protenix output.\"\"\"\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            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    f = data[\"input_feature_dict\"]\n    if \"centre_atom_mask\" in f:\n        return (f[\"centre_atom_mask\"] == 1).bool()\n    if \"center_atom_mask\" in f:\n        return (f[\"center_atom_mask\"] == 1).bool()\n\n    n_tokens = data.get(\"N_token\", torch.tensor(0)).item()\n    mask11 = (f[\"atom_to_tokatom_idx\"] == 11).bool()\n    mask12 = (f[\"atom_to_tokatom_idx\"] == 12).bool()\n    c11 = mask11.sum().item()\n    c12 = mask12.sum().item()\n    return mask11 if abs(c11 - n_tokens) < abs(c12 - n_tokens) else mask12\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:55.958851Z","iopub.execute_input":"2026-02-20T12:02:55.959066Z","iopub.status.idle":"2026-02-20T12:02:55.970485Z","shell.execute_reply.started":"2026-02-20T12:02:55.959048Z","shell.execute_reply":"2026-02-20T12:02:55.96982Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ─────────────── 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\"))\n\n# ── IMPROVEMENT 1: Tightened TBM thresholds for higher-quality templates ──\nMIN_SIMILARITY       = 0.0   # keep broad recall; filter by pct_id below\nMIN_PERCENT_IDENTITY = 60.0  # raised from 50 → better template quality\n\n# ── IMPROVEMENT 2: More Protenix recycling passes → better model refinement ──\nN_RECYCLE = int(os.environ.get(\"N_RECYCLE\", \"6\"))   # was implicitly 3; 6 gives meaningful gains\n\n# ── IMPROVEMENT 3: Confidence-weighted ensemble instead of first-come-first-served ──\nUSE_CONFIDENCE_WEIGHTING = True\n\nUSE_PROTENIX = True\n\n\ndef seed_everything(seed: int) -> None:\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    np.random.seed(seed)\n    torch.backends.cudnn.deterministic = 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    if os.path.isdir(\"/kaggle/working/protenix_data\"):\n        root_dir = \"/kaggle/working/protenix_data\"\n    elif os.path.isdir(\"/kaggle/working/protenix_root\"):\n        root_dir = \"/kaggle/working/protenix_root\"\n    else:\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 get_refined_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\"].bool()\n    atom_idx = f.get(\"atom_to_tokatom_idx\", None)\n    if atom_idx is not None:\n        c12 = (atom_idx == 12).sum().item()\n        n_tokens = data.get(\"N_token\", torch.tensor(0)).item()\n        if abs(c12 - n_tokens) <= 1:\n            return (atom_idx == 12).bool()\n        return (atom_idx == 11).bool()\n    return torch.zeros(f[\"atom_to_tokatom_idx\"].shape, dtype=torch.bool)\n\n\n# ─────────────── Geometry Utilities ──────────────────────────────────────────\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()\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)\n        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    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)\n        disp = rng.normal(0, amp, (6, 3))\n        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 adaptive_rna_constraints(coords, target_id, segments_map, confidence=1.0, passes=2) -> np.ndarray:\n    \"\"\"Enforce approximate C1'-C1' bond lengths (~5.9 Å) via gradient steps.\"\"\"\n    X = coords.copy()\n    segments = segments_map.get(target_id, [(0, len(X))])\n    strength = max(0.75 * (1.0 - min(confidence, 0.90)), 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: continue\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    return X\n\n\n# ── IMPROVEMENT 4: Smarter diversity — add global rotation + local torsion ──\ndef global_rotation_perturb(coords, rng, max_deg=30):\n    \"\"\"Apply a global random rotation around the molecular centroid.\"\"\"\n    X = coords.copy()\n    centroid = X.mean(0)\n    R = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-max_deg, max_deg))))\n    X = (X - centroid) @ R.T + centroid\n    return X\n\n\ndef local_torsion_wiggle(coords, segs, rng, n_pivots=4, deg=20):\n    \"\"\"Apply local torsion-like rotations at multiple random backbone pivots.\"\"\"\n    X = coords.copy()\n    for s, e in segs:\n        L = e - s\n        if L < 20: continue\n        pivots = rng.choice(np.arange(s + 5, e - 5), size=min(n_pivots, max(1, L//10)), replace=False)\n        for pivot in pivots:\n            R = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-deg, deg))))\n            p0 = X[pivot].copy()\n            X[pivot+1:e] = (X[pivot+1:e] - p0) @ R.T + p0\n    return X\n\n\ndef generate_diverse_ensemble(base_coords, n_samples, target_id, segments_map, confidence=0.7):\n    \"\"\"\n    IMPROVED: 5-strategy diversity pool instead of 3:\n      0 = hinge bend\n      1 = chain jitter\n      2 = smooth wiggle\n      3 = global rotation (NEW)\n      4 = local torsion wiggle (NEW)\n    Also applies constraint relaxation after each move.\n    \"\"\"\n    ensemble = [base_coords]\n    rng = np.random.default_rng(abs(hash(target_id)) % (2**32))\n    segs = segments_map.get(target_id, [(0, len(base_coords))])\n    while len(ensemble) < n_samples:\n        X = base_coords.copy()\n        choice = rng.integers(0, 5)\n        if choice == 0:\n            X = apply_hinge(X, segs[rng.integers(len(segs))], rng, deg=15)\n        elif choice == 1:\n            X = jitter_chains(X, segs, rng, trans=0.8)\n        elif choice == 2:\n            X = smooth_wiggle(X, segs, rng, amp=0.5)\n        elif choice == 3:\n            X = global_rotation_perturb(X, rng, max_deg=25)   # NEW\n        else:\n            X = local_torsion_wiggle(X, segs, rng)             # NEW\n        ensemble.append(adaptive_rna_constraints(X, target_id, segments_map, confidence=confidence))\n    return ensemble\n\n\n# ─────────────── TBM & Alignment ─────────────────────────────────────────────\ndef _make_aligner():\n    al = PairwiseAligner()\n    al.mode = \"global\"; al.match_score = 2; al.mismatch_score = -1.5\n    return al\n\n_aligner = _make_aligner()\n\n\ndef build_segments_map(df: pd.DataFrame):\n    seg_map = {}\n    for _, r in df.iterrows():\n        seq = r[\"sequence\"]\n        seg_map[r[\"target_id\"]] = [(0, len(seq))]\n    return seg_map, {}\n\n\ndef find_similar_sequences_detailed(query_seq, train_seqs_df, train_coords_dict, top_n=30):\n    \"\"\"IMPROVEMENT 5: Return similarity score alongside results for weighting.\"\"\"\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        # Quick pre-filter: length ratio\n        if abs(len(tseq) - len(query_seq)) / max(len(tseq), 1) > 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        results.append((tid, tseq, norm_s, train_coords_dict[tid], pct_id))\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):\n    \"\"\"Map template C1' coords onto query via pairwise alignment.\"\"\"\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            new_coords[i] = [i * 3, 0, 0]\n    return np.nan_to_num(new_coords)\n\n\ndef tbm_phase(test_df, train_seqs_df, train_coords_dict, segments_map):\n    \"\"\"\n    IMPROVEMENT 6: Sort TBM templates by pct_id (best first), then use\n    confidence-weighted diversity for fill-up instead of flat diversity.\n    \"\"\"\n    template_predictions, template_confidences, protenix_queue = {}, {}, {}\n    for _, row in test_df.iterrows():\n        tid, seq = row[\"target_id\"], row[\"sequence\"]\n        similar = find_similar_sequences_detailed(seq, train_seqs_df, train_coords_dict)\n        preds, confs = [], []\n        for tmpl_id, tmpl_seq, sim, tmpl_coords, pct_id in similar:\n            if len(preds) >= N_SAMPLE:\n                break\n            if sim < MIN_SIMILARITY or pct_id < MIN_PERCENT_IDENTITY:\n                continue\n            adapted = adapt_template_to_query(seq, tmpl_seq, tmpl_coords)\n            relaxed = adaptive_rna_constraints(adapted, tid, segments_map, confidence=sim)\n            preds.append(relaxed)\n            confs.append(pct_id / 100.0)   # normalise to [0,1]\n\n        template_predictions[tid] = preds\n        template_confidences[tid] = confs\n        if len(preds) < N_SAMPLE:\n            protenix_queue[tid] = (N_SAMPLE - len(preds), seq)\n    return template_predictions, template_confidences, protenix_queue\n\n\ndef coords_to_rows(target_id, seq, coords):\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            x, y, z = coords[s, i] if s < coords.shape[0] else (0, 0, 0)\n            row[f\"x_{s+1}\"], row[f\"y_{s+1}\"], row[f\"z_{s+1}\"] = float(x), float(y), float(z)\n        rows.append(row)\n    return rows\n\n\ndef generate_rna_structure(sequence, seed=None):\n    \"\"\"Helical de-novo fallback (unchanged).\"\"\"\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:55.982704Z","iopub.execute_input":"2026-02-20T12:02:55.98321Z","iopub.status.idle":"2026-02-20T12:02:56.018045Z","shell.execute_reply.started":"2026-02-20T12:02:55.98319Z","shell.execute_reply":"2026-02-20T12:02:56.017303Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ─────────────────────────────────────────────────────────────────────────────\n# Setup writable directory for Protenix templates / checkpoints\n# ─────────────────────────────────────────────────────────────────────────────\nNEW_ROOT = \"/kaggle/working/protenix_root\"\nos.makedirs(NEW_ROOT, exist_ok=True)\nos.environ[\"PROTENIX_ROOT_DIR\"] = NEW_ROOT\n\nsrc_pdb = \"/kaggle/input/stanford-rna-3d-folding-2/PDB_RNA\"\ndst_pdb = f\"{NEW_ROOT}/mmcif\"\nif not os.path.exists(dst_pdb):\n    try:\n        os.symlink(src_pdb, dst_pdb)\n        print(f\"Symlink: {dst_pdb} → {src_pdb}\")\n    except OSError:\n        import shutil\n        shutil.copytree(src_pdb, dst_pdb)\n\nORIGINAL_DATA_DIR = (\n    \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted\"\n    \"/Protenix-v1-adjust-v2/Protenix-v1-adjust-v2/Protenix-v1\"\n)\nfor folder in [\"common\", \"checkpoint\"]:\n    src = f\"{ORIGINAL_DATA_DIR}/{folder}\"\n    dst = f\"{NEW_ROOT}/{folder}\"\n    if not os.path.exists(dst):\n        if os.path.exists(src):\n            os.symlink(src, dst)\n            print(f\"Symlink: {dst} → {src}\")\n        else:\n            print(f\"WARNING: {folder} not found at {src}\")\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   = NEW_ROOT\n    return test_csv, output_csv, code_dir, root_dir\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:56.019341Z","iopub.execute_input":"2026-02-20T12:02:56.019591Z","iopub.status.idle":"2026-02-20T12:02:56.038978Z","shell.execute_reply.started":"2026-02-20T12:02:56.019572Z","shell.execute_reply":"2026-02-20T12:02:56.038283Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"WRITABLE_ROOT = \"/kaggle/working/protenix_data\"\nos.makedirs(WRITABLE_ROOT, exist_ok=True)\nos.environ[\"PROTENIX_ROOT_DIR\"] = WRITABLE_ROOT\n\nmmcif_target = f\"{WRITABLE_ROOT}/mmcif\"\nif not os.path.exists(mmcif_target):\n    source_pdb = \"/kaggle/input/stanford-rna-3d-folding-2/PDB_RNA\"\n    os.symlink(source_pdb, mmcif_target)\n    print(f\"Symlinked {source_pdb} → {mmcif_target}\")\n\nREADONLY_ASSETS = (\n    \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted\"\n    \"/Protenix-v1-adjust-v2/Protenix-v1-adjust-v2\"\n)\nfor folder in [\"common\", \"checkpoint\"]:\n    src = f\"{READONLY_ASSETS}/{folder}\"\n    dst = f\"{WRITABLE_ROOT}/{folder}\"\n    if not os.path.exists(dst) and os.path.exists(src):\n        os.symlink(src, dst)\n        print(f\"Symlinked {src} → {dst}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:56.039867Z","iopub.execute_input":"2026-02-20T12:02:56.040101Z","iopub.status.idle":"2026-02-20T12:02:56.045637Z","shell.execute_reply.started":"2026-02-20T12:02:56.040056Z","shell.execute_reply":"2026-02-20T12:02:56.045077Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Main Inference Pipeline\nKey improvements over baseline:\n1. **Higher TBM threshold** (`MIN_PERCENT_IDENTITY = 60`) → fewer but better templates\n2. **Increased Protenix recycling** (`N_RECYCLE = 6`) → more refined structures\n3. **Confidence-weighted ensemble** → best structures ranked first\n4. **5-strategy diversity** (added global rotation + local torsion) → better coverage\n5. **Template quality sorting** by percent identity before use\n","metadata":{}},{"cell_type":"code","source":"def main():\n    test_csv, output_csv, code_dir, root_dir = resolve_paths()\n    sys.path.insert(0, code_dir)\n    seed_everything(SEED)\n\n    test_df_full = pd.read_csv(test_csv)\n    test_df = (test_df_full.head(LOCAL_N_SAMPLES) if not IS_KAGGLE else test_df_full).reset_index(drop=True)\n\n    # ── Load TBM Pool ─────────────────────────────────────────────────────────\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    combined_seqs = pd.concat([train_seqs, val_seqs], ignore_index=True)\n    all_lbls      = pd.concat([train_labels, val_labels], ignore_index=True)\n\n    def process_labels(df):\n        coords = {}\n        for prefix, grp in df.groupby(df[\"ID\"].str.rsplit(\"_\", n=1).str[0]):\n            coords[prefix] = grp.sort_values(\"resid\")[[\"x_1\", \"y_1\", \"z_1\"]].values\n        return coords\n\n    train_coords = process_labels(all_lbls)\n    segments_map, _ = build_segments_map(test_df)\n\n    # ── Phase 1: Template-Based Modelling ────────────────────────────────────\n    template_preds, template_confs, protenix_queue = tbm_phase(\n        test_df, combined_seqs, train_coords, segments_map\n    )\n\n    # ── Phase 2: Protenix Deep-Learning Inference ────────────────────────────\n    protenix_preds, protenix_confs = {}, {}\n    if protenix_queue and USE_PROTENIX:\n        try:\n            from runner.inference import (\n                InferenceRunner, update_gpu_compatible_configs, update_inference_configs\n            )\n            from protenix.data.inference.infer_dataloader import InferenceDataset\n            from protenix.config.config import parse_configs\n\n            input_json_path = \"/kaggle/working/queue.json\"\n            q_list = [\n                {\n                    \"name\": k,\n                    \"sequences\": [{\"rnaSequence\": {\"sequence\": v[1], \"count\": 1}}]\n                }\n                for k, v in protenix_queue.items()\n            ]\n            with open(input_json_path, \"w\") as jf:\n                json.dump(q_list, jf)\n\n            # ── IMPROVEMENT 2: pass N_RECYCLE into configs ────────────────\n            arg_str = (\n                f\"--model_name {MODEL_NAME} \"\n                f\"--input_json_path {input_json_path} \"\n                f\"--dump_dir /kaggle/working/out \"\n                f\"--num_recycles {N_RECYCLE}\"\n            )\n            configs = parse_configs(arg_str=arg_str, fill_required_with_null=True)\n            configs = update_gpu_compatible_configs(configs)\n\n            # Attempt to set recycling on the config object directly (API varies)\n            for attr in [\"num_recycles\", \"n_recycles\", \"recycling_steps\"]:\n                if hasattr(configs, attr):\n                    setattr(configs, attr, N_RECYCLE)\n\n            runner  = InferenceRunner(configs)\n            dataset = InferenceDataset(configs)\n\n            for i in tqdm(range(len(dataset)), desc=\"Protenix inference\"):\n                data, _, _ = dataset[i]\n                tid = data.get(\"sample_name\")\n                if tid not in protenix_queue:\n                    continue\n                try:\n                    runner.update_model_configs(\n                        update_inference_configs(configs, data[\"N_token\"].item())\n                    )\n                    prediction = runner.predict(data)\n                    raw_coords = prediction[\"coordinate\"]\n\n                    # ── IMPROVEMENT 3: extract per-sample pLDDT confidence if available ──\n                    conf_score = 0.7  # default\n                    if \"plddt\" in prediction:\n                        plddt = prediction[\"plddt\"]\n                        mask  = get_refined_c1_mask(data).to(plddt.device)\n                        conf_score = float(plddt[:, mask].mean())\n\n                    mask   = get_refined_c1_mask(data).to(raw_coords.device)\n                    coords = raw_coords[:, mask, :].detach().cpu().numpy()\n                    protenix_preds[tid]  = coords\n                    protenix_confs[tid]  = conf_score\n                except Exception as e:\n                    print(f\"  Protenix failed for {tid}: {e}\")\n                finally:\n                    torch.cuda.empty_cache(); gc.collect()\n\n        except ImportError as e:\n            print(f\"Protenix not available: {e}; falling back to TBM + de-novo.\")\n\n    # ── Phase 3: Combine predictions ─────────────────────────────────────────\n    all_rows = []\n    for _, row in test_df.iterrows():\n        tid, seq = row[\"target_id\"], row[\"sequence\"]\n\n        combined   = list(template_preds.get(tid, []))\n        confs      = list(template_confs.get(tid, []))\n\n        # Fill with Protenix predictions (they carry higher confidence)\n        ptx = protenix_preds.get(tid)\n        if ptx is not None:\n            ptx_conf = protenix_confs.get(tid, 0.7)\n            for j in range(ptx.shape[0]):\n                if len(combined) >= N_SAMPLE:\n                    break\n                combined.append(ptx[j])\n                confs.append(ptx_conf)\n\n        # ── IMPROVEMENT 4: confidence-ranked diversity fill ───────────────\n        # Use the best available structure (highest conf) as the diversity seed\n        if 0 < len(combined) < N_SAMPLE:\n            best_idx  = int(np.argmax(confs)) if confs else 0\n            best_seed = combined[best_idx]\n            best_conf = confs[best_idx] if confs else 0.5\n            still_need = N_SAMPLE - len(combined)\n            variations = generate_diverse_ensemble(\n                best_seed, still_need + 1, tid, segments_map, confidence=best_conf\n            )\n            combined.extend(variations[1:still_need + 1])\n\n        # Ultimate fallback: de-novo helical structure\n        while len(combined) < N_SAMPLE:\n            dn = generate_rna_structure(seq, seed=abs(hash(tid)) % 10000 + len(combined))\n            combined.append(adaptive_rna_constraints(dn, tid, segments_map, confidence=0.1))\n\n        # ── IMPROVEMENT 3: sort final ensemble by confidence (descending) ──\n        if confs and len(confs) == len(combined):\n            order   = np.argsort(confs)[::-1]\n            combined = [combined[i] for i in order]\n\n        all_rows.extend(coords_to_rows(tid, seq, np.stack(combined[:N_SAMPLE])))\n\n    pd.DataFrame(all_rows).to_csv(output_csv, index=False)\n    print(f\"✓ Saved {len(all_rows)} residue rows to {output_csv}\")\n\n\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:02:56.047195Z","iopub.execute_input":"2026-02-20T12:02:56.047479Z","iopub.status.idle":"2026-02-20T12:03:28.177465Z","shell.execute_reply.started":"2026-02-20T12:02:56.047459Z","shell.execute_reply":"2026-02-20T12:03:28.17686Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Sanity check — view head of submission\nsubmission_path = \"/kaggle/working/submission.csv\"\nif os.path.exists(submission_path):\n    submission_df = pd.read_csv(submission_path)\n    print(f\"Shape: {submission_df.shape}\")\n    print(submission_df.head(10))\nelse:\n    print(\"submission.csv not yet generated — run main() first.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-02-20T12:03:28.17877Z","iopub.execute_input":"2026-02-20T12:03:28.179003Z","iopub.status.idle":"2026-02-20T12:03:28.192569Z","shell.execute_reply.started":"2026-02-20T12:03:28.178981Z","shell.execute_reply":"2026-02-20T12:03:28.191889Z"}},"outputs":[],"execution_count":null}]}