{"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python","version":"3.10.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"datasetVersion","sourceId":10880419,"datasetId":6760509,"databundleVersionId":11247150},{"sourceType":"datasetVersion","sourceId":10880374,"datasetId":6760482,"databundleVersionId":11247092},{"sourceType":"datasetVersion","sourceId":11899194,"datasetId":7479946,"databundleVersionId":12404228},{"sourceType":"datasetVersion","sourceId":11451236,"datasetId":7174725,"databundleVersionId":11892125},{"sourceType":"datasetVersion","sourceId":11230242,"datasetId":7014687,"databundleVersionId":11640305},{"sourceType":"datasetVersion","sourceId":13282339,"datasetId":7162026,"databundleVersionId":13982667},{"sourceType":"datasetVersion","sourceId":14874339,"datasetId":9502242,"databundleVersionId":15736806},{"sourceType":"datasetVersion","sourceId":11789546,"datasetId":7391536,"databundleVersionId":12277110},{"sourceType":"datasetVersion","sourceId":12615519,"datasetId":7943852,"databundleVersionId":13214987},{"sourceType":"datasetVersion","sourceId":14670281,"datasetId":9372180,"databundleVersionId":15512445},{"sourceType":"kernelVersion","sourceId":298071072},{"sourceType":"kernelVersion","sourceId":299151866},{"sourceType":"kernelVersion","sourceId":300520446},{"sourceType":"kernelVersion","sourceId":304615426}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Stanford RNA 3D Folding Part 2 — Conservative Rescue v5\n\nThis notebook is a **top-50 attempt** built around the strongest stable ideas from your working fallback path.\n\nWhat changed relative to the lower-scoring tuned versions:\n- **Kabsch-aligned template mixing** before any blend or swap\n- **Support-mask-aware consensus and patchwork** candidates\n- **Relaxed rescue template search** for weak-template targets\n- **Conservative heuristic selector** instead of validation-trained tuning\n- **Graceful Protenix auto-detect** when the offline runtime bundle is attached\n- **Hard coordinate sanitization** before writing `submission.csv`\n\nThe goal is to beat the unstable 0.386 / 0.369 branches without repeating the overfitting mistake.\n","metadata":{}},{"cell_type":"code","source":"import sys\nimport subprocess\nfrom pathlib import Path\n\ndef _wheel_compatible(path: Path) -> bool:\n    name = path.name.lower()\n    py_tag = f\"cp{sys.version_info.major}{sys.version_info.minor}\"\n    if name.endswith(\"none-any.whl\"):\n        return True\n    if py_tag in name:\n        return True\n    if \"abi3\" in name and \"linux\" in name:\n        return True\n    if \"py3-none\" in name or \"py2.py3-none\" in name:\n        return True\n    return False\n\ndef _find_best_wheel(keywords):\n    root = Path(\"/kaggle/input\")\n    if not root.exists():\n        return None\n    wheels = []\n    for p in root.rglob(\"*.whl\"):\n        name = p.name.lower()\n        if all(k in name for k in keywords):\n            wheels.append(p)\n    if not wheels:\n        return None\n\n    def score(p: Path):\n        name = p.name.lower()\n        py_tag = f\"cp{sys.version_info.major}{sys.version_info.minor}\"\n        return (\n            py_tag in name,\n            \"abi3\" in name,\n            \"manylinux\" in name,\n            -len(str(p)),\n            name\n        )\n    wheels.sort(key=score, reverse=True)\n    return wheels[0]\n\ndef _maybe_install(label, import_name, keywords):\n    try:\n        __import__(import_name)\n        print(f\"{label}: already available\")\n        return\n    except Exception:\n        pass\n\n    wheel = _find_best_wheel(keywords)\n    if wheel is None:\n        print(f\"{label}: wheel not found under /kaggle/input (skipping)\")\n        return\n    if not _wheel_compatible(wheel):\n        print(f\"{label}: found wheel but it looks incompatible with Python {sys.version_info.major}.{sys.version_info.minor}: {wheel.name}\")\n        return\n\n    cmd = [sys.executable, \"-m\", \"pip\", \"install\", \"--no-index\", \"--no-deps\", str(wheel)]\n    print(\"+\", \" \".join(cmd))\n    subprocess.run(cmd, check=False)\n\n_maybe_install(\"Biopython\", \"Bio\", [\"biopython\"])\n_maybe_install(\"biotite\", \"biotite\", [\"biotite\"])\n_maybe_install(\"RDKit\", \"rdkit\", [\"rdkit\"])\n_maybe_install(\"ml_collections\", \"ml_collections\", [\"ml\", \"collections\"])\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Notes\n\n- Designed for **Kaggle notebook submission runtime**\n- Auto-searches `/kaggle/input` for:\n  - competition CSVs\n  - wheel files\n  - Protenix repo\n  - Protenix offline checkpoint/cache bundle\n- If the full Protenix runtime is missing, it automatically falls back to the **conservative template-rescue** path\n- Output is written to `/kaggle/working/submission.csv` unless `SUBMISSION_CSV` is overridden\n- Safe defaults include `LAYERNORM_TYPE=torch` and `RNA_MSA_DEPTH_LIMIT=512` to avoid common runtime issues\n","metadata":{}},{"cell_type":"code","source":"import gc\nimport json\nimport math\nimport os\nimport random\nimport sys\nimport time\nimport shutil\nimport traceback\nimport warnings\nfrom pathlib import Path\n\nos.environ.setdefault(\"LAYERNORM_TYPE\", \"torch\")\nos.environ.setdefault(\"RNA_MSA_DEPTH_LIMIT\", \"512\")\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom tqdm import tqdm\n\nwarnings.filterwarnings(\"ignore\")\n\n# ───────────────────────── Paths & Constants ─────────────────────────────\nCOMP_DATA_DIR = \"/kaggle/input/competitions/stanford-rna-3d-folding-2\"\nDEFAULT_TEST_CSV = f\"{COMP_DATA_DIR}/test_sequences.csv\"\nDEFAULT_TRAIN_CSV = f\"{COMP_DATA_DIR}/train_sequences.csv\"\nDEFAULT_TRAIN_LBLS = f\"{COMP_DATA_DIR}/train_labels.csv\"\nDEFAULT_VAL_CSV = f\"{COMP_DATA_DIR}/validation_sequences.csv\"\nDEFAULT_VAL_LBLS = f\"{COMP_DATA_DIR}/validation_labels.csv\"\nDEFAULT_OUTPUT = \"/kaggle/working/submission.csv\"\n\nMODEL_NAME = os.environ.get(\"MODEL_NAME\", \"protenix_base_20250630_v1.0.0\")\nDEFAULT_CODE_DIR = os.environ.get(\n    \"PROTENIX_CODE_DIR_DEFAULT\",\n    \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted/Protenix-v1-adjust-v2/Protenix-v1-adjust-v2/Protenix-v1\",\n)\nDEFAULT_ROOT_DIR = os.environ.get(\"PROTENIX_ROOT_DIR_DEFAULT\", DEFAULT_CODE_DIR)\n\n# --- PROMPT 6: N_SAMPLE = 7 ---\nN_SAMPLE = 7\nSEED = int(os.environ.get(\"SEED\", \"42\"))\nMAX_SEQ_LEN = int(os.environ.get(\"MAX_SEQ_LEN\", \"512\"))\nCHUNK_OVERLAP = int(os.environ.get(\"CHUNK_OVERLAP\", \"128\"))\nMAX_LONG_SAMPLES = int(os.environ.get(\"MAX_LONG_SAMPLES\", \"3\"))\n\n# --- PROMPT 1: THRESHOLDS ---\nMAX_TEMPLATE_SEARCH = 12\nMAX_TEMPLATE_CANDIDATES = 12\nTEMPLATE_ONLY_CONF = 0.78\nTEMPLATE_STRONG_CONF = 0.65\nTEMPLATE_MEDIUM_CONF = 0.58\nTEMPLATE_WEAK_CONF = 0.50\nMIN_TEMPLATE_QUALITY_KEEP = 0.32\n\nQUALITY_DELTA = 0.10\n# --- PROMPT 4: DIVERSITY ---\nDIVERSITY_LAMBDA = 0.045\nDIVERSITY_CAP = 4.0\nRMSD_SUBSAMPLE = 500\nAUGMENT_TEMPLATE_VARIANTS = 3\n\nSHORT_PTX_BASE_QUALITY = 0.67\nLONG_PTX_BASE_QUALITY = 0.61\n\nCOORD_CLIP_MIN = -999.999\nCOORD_CLIP_MAX = 9999.999\nABSURD_COORD_THRESHOLD = 100000.0\n\nUSE_PROTENIX = str(os.environ.get(\"USE_PROTENIX\", \"true\")).strip().lower() not in {\"0\", \"false\", \"no\", \"off\"}\nIS_KAGGLE = os.path.exists(\"/kaggle/input\")\nLOCAL_N_SAMPLES = None if IS_KAGGLE else 2\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\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\n# [CÁC HÀM TIỆN ÍCH DÀI GIỮ NGUYÊN...]\ndef seed_everything(seed: int) -> None:\n    os.environ[\"CUBLAS_WORKSPACE_CONFIG\"] = \":4096:8\"\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    if torch.cuda.is_available():\n        torch.cuda.manual_seed(seed)\n        torch.cuda.manual_seed_all(seed)\n    torch.backends.cudnn.benchmark = False\n    torch.backends.cudnn.deterministic = True\n\ndef _auto_find_competition_data() -> str | None:\n    root = Path(\"/kaggle/input\")\n    if not root.exists(): return None\n    candidates = [p for p in root.rglob(\"*\") if p.is_dir() and {\"train_sequences.csv\", \"test_sequences.csv\"}.issubset({x.name for x in p.iterdir() if x.is_file()})]\n    if not candidates: return None\n    candidates.sort(key=lambda p: (\"stanford-rna-3d-folding-2\" not in str(p).lower(), len(str(p))))\n    return str(candidates[0])\n\ndef _walk_input_dirs():\n    root = Path(\"/kaggle/input\")\n    if root.exists():\n        for dirpath, dirnames, filenames in os.walk(root): yield Path(dirpath)\n\ndef _path_has_protenix_repo(path: str | Path | None) -> bool:\n    if not path: return False\n    p = Path(path)\n    return p.is_dir() and (p / \"runner\" / \"inference.py\").exists() and ((p / \"configs\" / \"configs_base.py\").exists() or (p / \"configs\").is_dir())\n\ndef _resolve_protenix_runtime_layout(path: str | Path | None, model_name: str):\n    if not path: return None\n    p = Path(path)\n    layouts = [\n        {\"kind\": \"common\", \"root\": p, \"checkpoint\": p / \"checkpoint\" / f\"{model_name}.pt\", \"ccd\": p / \"common\" / \"components.cif\", \"ccd_cache\": p / \"common\" / \"components.cif.rdkit_mol.pkl\"},\n        {\"kind\": \"release_data\", \"root\": p, \"checkpoint\": p / \"release_data\" / \"checkpoint\" / f\"{model_name}.pt\", \"ccd\": p / \"release_data\" / \"ccd_cache\" / \"components.cif\", \"ccd_cache\": p / \"release_data\" / \"ccd_cache\" / \"components.cif.rdkit_mol.pkl\"}\n    ]\n    for layout in layouts:\n        if layout[\"checkpoint\"].exists() and layout[\"ccd\"].exists() and layout[\"ccd_cache\"].exists(): return layout\n    return None\n\ndef _path_has_protenix_runtime(path: str | Path | None, model_name: str) -> bool:\n    return _resolve_protenix_runtime_layout(path, model_name) is not None\n\ndef _safe_link_or_copy(src: Path | str | None, dst: Path) -> None:\n    if src is None or not Path(src).exists() or dst.exists(): return\n    src = Path(src)\n    dst.parent.mkdir(parents=True, exist_ok=True)\n    try: os.symlink(src, dst)\n    except: shutil.copytree(src, dst, dirs_exist_ok=True) if src.is_dir() else shutil.copy2(src, dst)\n\ndef _materialize_runtime_root(layout: dict, model_name: str) -> str:\n    if layout[\"kind\"] == \"common\": return str(layout[\"root\"])\n    norm_root = Path(\"/kaggle/working/protenix_runtime_auto\")\n    _safe_link_or_copy(layout[\"checkpoint\"], norm_root / \"checkpoint\" / f\"{model_name}.pt\")\n    _safe_link_or_copy(layout[\"ccd\"], norm_root / \"common\" / \"components.cif\")\n    _safe_link_or_copy(layout[\"ccd_cache\"], norm_root / \"common\" / \"components.cif.rdkit_mol.pkl\")\n    return str(norm_root)\n\ndef _finalize_protenix_paths(code_dir, root_dir, model_name):\n    if not _path_has_protenix_runtime(root_dir, model_name):\n        cands = [p for p in _walk_input_dirs() or [] if _path_has_protenix_runtime(p, model_name)]\n        if cands: root_dir = str(cands[0])\n    if not _path_has_protenix_repo(code_dir):\n        cands = [p for p in _walk_input_dirs() or [] if _path_has_protenix_repo(p)]\n        if cands: code_dir = str(cands[0])\n    return code_dir, root_dir\n\ndef resolve_paths():\n    data_dir = _auto_find_competition_data() or COMP_DATA_DIR\n    test_csv = str(Path(data_dir) / \"test_sequences.csv\")\n    train_csv = str(Path(data_dir) / \"train_sequences.csv\")\n    train_lbls = str(Path(data_dir) / \"train_labels.csv\")\n    val_csv = str(Path(data_dir) / \"validation_sequences.csv\")\n    val_lbls = str(Path(data_dir) / \"validation_labels.csv\")\n    code_dir, root_dir = _finalize_protenix_paths(DEFAULT_CODE_DIR, DEFAULT_ROOT_DIR, MODEL_NAME)\n    return test_csv, train_csv, train_lbls, val_csv, val_lbls, DEFAULT_OUTPUT, code_dir, root_dir\n\ndef ensure_required_files(root_dir: str) -> str:\n    layout = _resolve_protenix_runtime_layout(root_dir, MODEL_NAME)\n    return _materialize_runtime_root(layout, MODEL_NAME)\n\ndef build_input_json(df: pd.DataFrame, json_path: str) -> None:\n    data = [{\"name\": row[\"target_id\"], \"covalent_bonds\": [], \"sequences\": [{\"rnaSequence\": {\"sequence\": row[\"sequence\"], \"count\": 1}}]} for _, row in df.iterrows()]\n    with open(json_path, \"w\", encoding=\"utf-8\") as f: json.dump(data, f)\n\ndef build_configs(input_json_path, dump_dir, model_name):\n    from configs.configs_base import configs as 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    cfg = {**base, **{\"data\": data_configs}, **inference_configs}\n    def deep_update(t, p):\n        for k, v in p.items(): deep_update(t[k], v) if isinstance(v, dict) and k in t and isinstance(t[k], dict) else t.update({k: v})\n    deep_update(cfg, model_configs[model_name])\n    arg_str = f\"--model_name {model_name} --input_json_path {input_json_path} --dump_dir {dump_dir} --sample_diffusion.N_sample {N_SAMPLE} --seeds {SEED}\"\n    return parse_configs(configs=cfg, arg_str=arg_str, fill_required_with_null=True)\n\ndef coords_to_rows(target_id: str, seq: str, coords: np.ndarray) -> list:\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] and i < coords.shape[1] else (0.0, 0.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\ndef split_into_chunks(seq_len: int, max_len: int, overlap: int):\n    if seq_len <= max_len: return [(0, seq_len)]\n    return [(p, min(p + max_len, seq_len)) for p in range(0, seq_len, max_len - overlap) if p < seq_len]\n\ndef kabsch_align(P: np.ndarray, Q: np.ndarray):\n    P, Q = np.asarray(P, dtype=float), np.asarray(Q, dtype=float)\n    cP, cQ = P.mean(axis=0), Q.mean(axis=0)\n    Pc, Qc = P - cP, Q - cQ\n    if len(P) < 3 or len(Q) < 3: return np.eye(3), cQ - cP\n    H = Pc.T @ Qc\n    U, _, Vt = np.linalg.svd(H)\n    S = np.eye(3)\n    if np.linalg.det(Vt.T @ U.T) < 0: S[2, 2] = -1\n    R = Vt.T @ S @ U.T\n    return R, cQ - R @ cP\n\n# --- PROMPT 6: KABSCH STITCHING ---\ndef stitch_chunk_coords(chunk_coords_list: list, chunk_ranges: list, seq_len: int) -> np.ndarray:\n    if len(chunk_coords_list) == 1:\n        coords = chunk_coords_list[0]\n        if coords.shape[0] >= seq_len: return coords[:seq_len]\n        out = np.zeros((seq_len, 3), dtype=coords.dtype)\n        out[:coords.shape[0]] = coords\n        return out\n\n    aligned = [chunk_coords_list[0].copy()]\n    for i in range(1, len(chunk_coords_list)):\n        prev_start, prev_end = chunk_ranges[i - 1]\n        cur_start, cur_end = chunk_ranges[i]\n        ov_start, ov_end = cur_start, min(prev_end, cur_end)\n        \n        if (ov_end - ov_start) < 3:\n            aligned.append(chunk_coords_list[i].copy())\n            continue\n            \n        prev_ov = aligned[i - 1][ov_start - prev_start: ov_end - prev_start]\n        cur_ov = chunk_coords_list[i][ov_start - cur_start: ov_end - cur_start]\n        valid = ~(np.isnan(prev_ov).any(axis=1) | np.isnan(cur_ov).any(axis=1))\n        \n        if valid.sum() < 3:\n            aligned.append(chunk_coords_list[i].copy())\n            continue\n            \n        R, t = kabsch_align(cur_ov[valid], prev_ov[valid])\n        aligned.append((chunk_coords_list[i] @ R.T) + t)\n        \n    full, weights = np.zeros((seq_len, 3), dtype=np.float64), np.zeros(seq_len, dtype=np.float64)\n    for i, ((s, e), coords) in enumerate(zip(chunk_ranges, aligned)):\n        actual_end = min(s + coords.shape[0], seq_len)\n        used_len = actual_end - s\n        w = np.ones(used_len, dtype=np.float64)\n        if i > 0:\n            ramp_len = min(chunk_ranges[i - 1][1], e) - s\n            if ramp_len > 0: w[:ramp_len] = np.linspace(0.0, 1.0, ramp_len)\n        if i < len(chunk_ranges) - 1:\n            next_s = chunk_ranges[i + 1][0]\n            ramp_len = actual_end - next_s\n            if ramp_len > 0 and (next_s - s) < used_len: w[(next_s - s):used_len] = np.linspace(1.0, 0.0, ramp_len)\n        full[s:actual_end] += coords[:used_len] * w[:, None]\n        weights[s:actual_end] += w\n    mask = weights > 0\n    full[mask] /= weights[mask, None]\n    return full\n\nclass _FallbackAligner:\n    def score(self, q, t):\n        ml = min(len(q), len(t))\n        if ml == 0: return 0.0\n        return 2.0 * ml * (min(len(q), len(t)) / max(len(q), len(t)))\n    def align(self, q, t): return type('obj', (object,), {'aligned': ([(0, min(len(q), len(t)))], [(0, min(len(q), len(t)))])})()\n\ndef make_aligner():\n    try:\n        from Bio.Align import PairwiseAligner\n        al = PairwiseAligner()\n        al.mode, al.match_score, al.mismatch_score, al.open_gap_score, al.extend_gap_score = \"global\", 2, -1.5, -8, -0.4\n        al.query_left_open_gap_score = al.target_left_open_gap_score = -8\n        al.query_left_extend_gap_score = al.target_left_extend_gap_score = -0.4\n        return al\n    except: return _FallbackAligner()\n\n_aligner = make_aligner()\n\ndef get_chain_segments(row) -> list:\n    return [(0, len(row[\"sequence\"]))]\n\ndef build_segments_map(df: pd.DataFrame) -> tuple[dict, dict]:\n    return {r[\"target_id\"]: get_chain_segments(r) for _, r in df.iterrows()}, {}\n\ndef process_labels(labels_df: pd.DataFrame) -> dict:\n    return {prefix: grp.sort_values(\"resid\")[[\"x_1\", \"y_1\", \"z_1\"]].values for prefix, grp in labels_df.groupby(labels_df[\"ID\"].str.rsplit(\"_\", n=1).str[0])}\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)): 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\ndef _estimate_direction_from_coords(coords: np.ndarray, idx: int) -> np.ndarray:\n    for rng in [range(idx - 1, -1, -1), range(idx + 1, len(coords))]:\n        valid = [j for j in rng if np.isfinite(coords[j]).all()]\n        if len(valid) >= 2:\n            v = coords[valid[0]] - coords[valid[1]]\n            n = np.linalg.norm(v)\n            if n > 1e-6: return v / n\n    return np.array([1.0, 0.0, 0.0], dtype=float)\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, dtype=float)\n    for (qs, qe), (ts, te) in zip(*aln.aligned):\n        if len(template_coords[ts:te]) == (qe - qs): new_coords[qs:qe] = template_coords[ts:te]\n    for i in range(len(new_coords)):\n        if np.isnan(new_coords[i, 0]): new_coords[i] = np.array([i * 5.95, 0.0, 0.0], dtype=float)\n    return np.nan_to_num(new_coords)\n\ndef adaptive_rna_constraints(coords, target_id, segments_map, confidence=1.0, passes=2) -> np.ndarray:\n    X = np.asarray(coords, dtype=float).copy()\n    strength = max(0.75 * (1.0 - min(confidence, 0.97)), 0.02)\n    for _ in range(passes):\n        for s, e in segments_map.get(target_id, [(0, len(X))]):\n            C = X[s:e]\n            if len(C) < 3: continue\n            d = C[1:] - C[:-1]\n            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            C[1:-1] += (0.06 * strength) * (0.5 * (C[:-2] + C[2:]) - C[1:-1])\n            X[s:e] = C\n    return X\n\ndef _rotmat(axis, ang):\n    a = np.asarray(axis, float)\n    a /= np.linalg.norm(a) + 1e-12\n    c, s, CC = np.cos(ang), np.sin(ang), 1 - np.cos(ang)\n    x, y, z = a\n    return np.array([[c+x*x*CC, x*y*CC-z*s, x*z*CC+y*s], [y*x*CC+z*s, c+y*y*CC, y*z*CC-x*s], [z*x*CC-y*s, z*y*CC+x*s, c+z*z*CC]])\n\ndef generate_rna_structure(sequence: str, seed=None) -> np.ndarray:\n    if seed: np.random.seed(seed)\n    return np.array([[10.0*np.cos(i*0.6), 10.0*np.sin(i*0.6), i*2.5] for i in range(len(sequence))], dtype=float)\n\n# --- PROMPT 5: SANITIZATION PASSES ---\ndef sanitize_coords(coords: np.ndarray, target_id: str, segments_map: dict,\n                    confidence: float = 0.5, clip: bool = True, source: str = \"protenix\") -> np.ndarray:\n    X = np.asarray(coords, dtype=float).copy()\n    if X.ndim != 2 or X.shape[1] != 3: X = X.reshape(-1, 3)\n    X[np.abs(X) > ABSURD_COORD_THRESHOLD] = np.nan\n    X[~np.isfinite(X)] = np.nan\n    for s, e in segments_map.get(target_id, [(0, len(X))]):\n        C = X[s:e].copy()\n        if not np.isfinite(C).all(axis=1).any():\n            X[s:e] = generate_rna_structure(\"A\" * len(C))\n            continue\n        X[s:e] = C\n    \n    n_passes = 1 if \"template\" in source else 2\n    X = adaptive_rna_constraints(X, target_id, segments_map, confidence=confidence, passes=n_passes)\n    if clip: X = np.clip(X, COORD_CLIP_MIN, COORD_CLIP_MAX)\n    return X\n\ndef sanitize_stack(pred_stack: np.ndarray, target_id: str, segments_map: dict, confidence: float = 0.5) -> np.ndarray:\n    return np.stack([sanitize_coords(p, target_id, segments_map, confidence=confidence, source=\"protenix\") for p in pred_stack], axis=0)\n\ndef clamp01(x): return float(max(0.0, min(1.0, x)))\n\n# --- PROMPT 2: WEIGHTS ---\ndef compute_template_features(query_seq, template_seq, similarity, percent_identity, aligned_query, aligned_template):\n    aligned_len = max(1, len(aligned_query))\n    matched = sum((q != \"-\" and t != \"-\") for q, t in zip(aligned_query, aligned_template))\n    coverage = matched / max(1, len(query_seq))\n    gap_fraction = 1.0 - (matched / aligned_len)\n    length_ratio = min(len(query_seq), len(template_seq)) / max(len(query_seq), len(template_seq))\n\n    support_mask = np.array([1.0 if q!=\"-\" and t!=\"-\" else 0.0 for q, t in zip(aligned_query, aligned_template) if q!=\"-\"])\n    identity_mask = np.array([1.0 if q==t and q!=\"-\" else 0.0 for q, t in zip(aligned_query, aligned_template) if q!=\"-\"])\n\n    base_conf = (\n        0.30 * clamp01(similarity) +\n        0.27 * clamp01(percent_identity / 100.0) +\n        0.28 * clamp01(coverage) +\n        0.13 * clamp01(length_ratio) -\n        0.08 * clamp01(gap_fraction)\n    )\n    confidence = clamp01(0.97 * base_conf + 0.02 * float(np.mean(identity_mask) if len(identity_mask) else 0) + 0.01 * float(np.mean(support_mask) if len(support_mask) else 0))\n\n    return {\"coverage\": coverage, \"gap_fraction\": gap_fraction, \"length_ratio\": length_ratio, \"confidence\": confidence, \"support_mask\": support_mask, \"identity_mask\": identity_mask}\n\ndef find_similar_sequences_detailed(query_seq, train_seqs_df, train_coords_dict, top_n=12):\n    similar = []\n    for _, r in train_seqs_df.iterrows():\n        if r[\"target_id\"] in train_coords_dict and abs(len(r[\"sequence\"]) - len(query_seq))/max(len(r[\"sequence\"]), max(1, len(query_seq))) <= 0.60:\n            similar.append((r[\"target_id\"], r[\"sequence\"], _aligner.score(query_seq, r[\"sequence\"]) / (2 * min(len(query_seq), len(r[\"sequence\"]))), train_coords_dict[r[\"target_id\"]]))\n    similar.sort(key=lambda x: x[2], reverse=True)\n    results = []\n    for tid, tseq, sim, coords in similar[:top_n*6]:\n        aln = next(iter(_aligner.align(query_seq, tseq)))\n        pct_id = 100 * sum(query_seq[qp] == tseq[tp] for (qs,qe), (ts,te) in zip(*aln.aligned) for qp,tp in zip(range(qs,qe), range(ts,te))) / max(1, len(query_seq))\n        aq, at = _build_aligned_strings(query_seq, tseq, aln)\n        results.append((tid, tseq, sim, coords, pct_id, aq, at))\n    return sorted(results, key=lambda x: x[2], reverse=True)[:top_n]\n\ndef build_template_prediction_candidates(sequence, target_id, train_seqs_df, train_coords_dict, segments_map, max_candidates=MAX_TEMPLATE_SEARCH):\n    candidates = []\n    for rank, (tmpl_id, tmpl_seq, sim, coords, pct_id, aq, at) in enumerate(find_similar_sequences_detailed(sequence, train_seqs_df, train_coords_dict, max_candidates), 1):\n        if tmpl_id == target_id: continue\n        feat = compute_template_features(sequence, tmpl_seq, sim, pct_id, aq, at)\n        if feat[\"confidence\"] < MIN_TEMPLATE_QUALITY_KEEP and pct_id < 20.0: continue\n        adapted = adapt_template_to_query(sequence, tmpl_seq, coords)\n        candidates.append({\n            \"pred_id\": f\"template::{tmpl_id}\", \"source\": \"template\", \"template_id\": tmpl_id, \"similarity\": float(sim), \"percent_identity\": float(pct_id),\n            \"quality\": float(feat[\"confidence\"]), \"coords\": sanitize_coords(adapted, target_id, segments_map, confidence=feat[\"confidence\"], source=\"template\"),\n            \"support_mask\": feat[\"support_mask\"], \"identity_mask\": feat[\"identity_mask\"], \"rank_hint\": rank\n        })\n    uniq = {c[\"template_id\"]: c for c in sorted(candidates, key=lambda x: x[\"quality\"], reverse=True)}\n    return list(uniq.values())[:max_candidates]\n\ndef route_from_template_candidates(candidate_meta, seq_len):\n    if not candidate_meta: return {\"route\": \"fallback_heavy\", \"template_slots\": 0, \"protenix_slots\": N_SAMPLE, \"best_confidence\": 0.0}\n    best_conf = candidate_meta[0][\"quality\"]\n    strong = sum(c[\"quality\"] >= TEMPLATE_STRONG_CONF for c in candidate_meta)\n    med = sum(c[\"quality\"] >= TEMPLATE_MEDIUM_CONF for c in candidate_meta)\n    if best_conf >= TEMPLATE_ONLY_CONF and strong >= 4: route, t, p = \"template_only\", N_SAMPLE, 0\n    elif best_conf >= TEMPLATE_STRONG_CONF and med >= 2: route, t, p = \"hybrid\", min(3, N_SAMPLE-1), max(1, N_SAMPLE-3)\n    elif best_conf >= TEMPLATE_MEDIUM_CONF: route, t, p = \"hybrid\", 2, N_SAMPLE-2\n    elif best_conf >= TEMPLATE_WEAK_CONF: route, t, p = \"fallback_heavy\", 1, N_SAMPLE-1\n    else: route, t, p = \"fallback_heavy\", 0, N_SAMPLE\n    if seq_len > 800 and t > 0: t, p = max(t, min(3, len(candidate_meta))), max(0, N_SAMPLE - t)\n    return {\"route\": route, \"template_slots\": t, \"protenix_slots\": p, \"best_confidence\": best_conf}\n\ndef make_template_aug_candidates(best_cand: dict, target_id: str, segments_map: dict, n_aug: int = 2, seed: int = 0) -> list[dict]:\n    if not best_cand or n_aug <= 0: return []\n    rng = np.random.default_rng(seed)\n    segs = segments_map.get(target_id, [(0, len(best_cand[\"coords\"]))])\n    out = []\n    for i in range(n_aug):\n        v = smooth_wiggle(best_cand[\"coords\"], segs, rng, amp=0.45) if i==0 else apply_hinge(best_cand[\"coords\"], max(segs, key=lambda s: s[1]-s[0]), rng, 14)\n        out.append({**best_cand, \"pred_id\": f\"template_aug::{i}\", \"source\": \"template_aug\", \"quality\": max(0.05, best_cand[\"quality\"] - 0.035*(i+1)),\n                    \"coords\": sanitize_coords(v, target_id, segments_map, confidence=best_cand[\"quality\"], source=\"template_aug\")})\n    return out\n\ndef make_prediction_aug_candidates(base_coords: np.ndarray, target_id: str, segments_map: dict, n_aug: int = 2, seed: int = 0, source_prefix=\"protenix_aug\", base_quality=0.55):\n    if n_aug <= 0: return []\n    rng = np.random.default_rng(seed)\n    segs = segments_map.get(target_id, [(0, len(base_coords))])\n    out = []\n    for i in range(n_aug):\n        v = smooth_wiggle(base_coords, segs, rng, amp=0.35) if i==0 else apply_hinge(base_coords, max(segs, key=lambda s: s[1]-s[0]), rng, 10)\n        out.append({\"pred_id\": f\"{source_prefix}::{i}\", \"source\": source_prefix, \"quality\": max(0.05, base_quality - 0.03*(i+1)),\n                    \"coords\": sanitize_coords(v, target_id, segments_map, confidence=base_quality, source=source_prefix)})\n    return out\n\ndef make_de_novo_candidate(sequence, target_id, segments_map, variant_idx=0, seed=0):\n    base = generate_rna_structure(sequence, seed=seed)\n    return {\"pred_id\": f\"denovo::{variant_idx}\", \"source\": \"de_novo\", \"quality\": max(0.02, 0.18 - 0.02*variant_idx),\n            \"coords\": sanitize_coords(base, target_id, segments_map, confidence=0.20, source=\"de_novo\")}\n\ndef aligned_rmsd(A, B, max_points=500):\n    A, B = np.asarray(A, dtype=float)[:min(len(A), len(B))], np.asarray(B, dtype=float)[:min(len(A), len(B))]\n    if len(A) == 0: return 0.0\n    if len(A) > max_points: A, B = A[np.linspace(0, len(A)-1, max_points).astype(int)], B[np.linspace(0, len(A)-1, max_points).astype(int)]\n    R, t = kabsch_align(A, B)\n    return float(np.sqrt(np.mean(np.sum(((A @ R.T) + t - B)**2, axis=1))))\n\n# --- PROMPT 4: HARD FILTER DIVERSITY ---\ndef select_final_five_scored(candidates, n_out=N_SAMPLE):\n    if not candidates: return []\n    cands = sorted([dict(c) for c in candidates], key=lambda x: x.get(\"quality\", 0.0), reverse=True)\n    selected = [cands[0]]\n    remaining = cands[1:]\n    while remaining and len(selected) < n_out:\n        best_val = -1e9; best_item = None\n        for cand in remaining:\n            val = float(cand.get(\"quality\", 0.0))\n            if selected:\n                div = min(aligned_rmsd(cand[\"coords\"], s[\"coords\"]) for s in selected)\n                if div < 1.5: continue  # LỌC CỨNG (HARD FILTER)\n                val += DIVERSITY_LAMBDA * min(div, DIVERSITY_CAP)\n            if val > best_val: best_val, best_item = val, cand\n        if best_item:\n            selected.append(best_item)\n            remaining.remove(best_item)\n        else: break\n    return selected[:n_out]\n\ndef plan_targets(test_df, train_seqs_df, train_coords_dict, segments_map):\n    print(f\"\\n{'=' * 72}\\nPHASE 1: template search + routing\\n{'=' * 72}\")\n    plans, p_queue = {}, {}\n    for r_idx, row in test_df.reset_index(drop=True).iterrows():\n        tid, seq = row[\"target_id\"], row[\"sequence\"]\n        tmpls = build_template_prediction_candidates(seq, tid, train_seqs_df, train_coords_dict, segments_map, MAX_TEMPLATE_CANDIDATES)\n        route = route_from_template_candidates(tmpls, len(seq))\n        plans[tid] = {\"target_id\": tid, \"sequence\": seq, \"template_candidates\": tmpls, \"route_info\": route, \"row_idx\": r_idx}\n        if route[\"protenix_slots\"] > 0: p_queue[tid] = {\"n_needed\": route[\"protenix_slots\"], \"sequence\": seq}\n        print(f\"  {tid:>10s} | best={tmpls[0]['template_id'] if tmpls else 'none'} conf={route['best_confidence']:.3f} | route={route['route']} t={route['template_slots']} p={route['protenix_slots']}\")\n    return plans, p_queue\n\ndef run_protenix_phase(protenix_queue: dict) -> dict:\n    if not protenix_queue or not USE_PROTENIX: return {}\n    print(f\"\\n{'=' * 72}\\nPHASE 2: Protenix for {len(protenix_queue)} routed targets\\n{'=' * 72}\")\n    work_dir = Path(\"/kaggle/working\")\n    work_dir.mkdir(parents=True, exist_ok=True)\n    tasks, chunk_info, infer_slots_map = [], {}, {}\n\n    for tid, info in protenix_queue.items():\n        seq, n_needed = info[\"sequence\"], int(info[\"n_needed\"])\n        # --- PROMPT 6: N_INFER = 7 ---\n        n_infer = max(n_needed + 2, 7) if len(seq) <= MAX_SEQ_LEN else min(n_needed, MAX_LONG_SAMPLES)\n        infer_slots_map[tid] = n_infer\n        if len(seq) <= MAX_SEQ_LEN:\n            tasks.append({\"target_id\": tid, \"sequence\": seq})\n            chunk_info[tid] = [{\"name\": tid, \"range\": (0, len(seq))}]\n        else:\n            chunks = split_into_chunks(len(seq), MAX_SEQ_LEN, CHUNK_OVERLAP)\n            chunk_info[tid] = []\n            for ci, (cs, ce) in enumerate(chunks):\n                cname = f\"{tid}_chunk{ci}\"\n                tasks.append({\"target_id\": cname, \"sequence\": seq[cs:ce]})\n                chunk_info[tid].append({\"name\": cname, \"range\": (cs, ce)})\n\n    if not tasks: return {}\n    input_json_path = str(work_dir / \"protenix_queue_input.json\")\n    build_input_json(pd.DataFrame(tasks), input_json_path)\n\n    # --- BÙA TRẤN YỂM LỖI BIOTITE ---\n    import sys\n    q_path = \"/kaggle/input/datasets/qiweiyin/protenix-v1-adjusted/Protenix-v1-adjust-v2/Protenix-v1-adjust-v2/Protenix-v1\"\n    if q_path not in sys.path: sys.path.insert(0, q_path)\n    try:\n        import biotite.structure.io.pdbx.convert as pdbx_convert\n        if not hasattr(pdbx_convert, \"PDBX_BOND_TYPE_ID_TO_TYPE\"): pdbx_convert.PDBX_BOND_TYPE_ID_TO_TYPE = {}\n    except Exception: pass\n    # ---------------------------------\n\n    from protenix.data.inference.infer_dataloader import InferenceDataset\n    from runner.inference import InferenceRunner, update_gpu_compatible_configs, update_inference_configs\n\n    configs = update_gpu_compatible_configs(build_configs(input_json_path, str(work_dir / \"outputs\"), MODEL_NAME))\n    runner = InferenceRunner(configs)\n    dataset = InferenceDataset(configs)\n    raw_predictions = {}\n\n    for i in tqdm(range(len(dataset)), desc=\"Protenix Inference\"):\n        data, atom_array, err = dataset[i]\n        s_name = data.get(\"sample_name\", f\"sample_{i}\")\n        if err:\n            raw_predictions[s_name] = None\n            continue\n        try:\n            tid = s_name.split(\"_chunk\")[0] if \"_chunk\" in s_name else s_name\n            new_cfg = update_inference_configs(configs, int(data[\"N_token\"].item()))\n            new_cfg.sample_diffusion.N_sample = infer_slots_map.get(tid, N_SAMPLE)\n            runner.update_model_configs(new_cfg)\n            raw_coords = runner.predict(data)[\"coordinate\"]\n            mask = (data[\"input_feature_dict\"].get(\"centre_atom_mask\", torch.ones(raw_coords.shape[1])) == 1).to(raw_coords.device)\n            raw_predictions[s_name] = raw_coords[:, mask, :].detach().cpu().numpy()\n        except: raw_predictions[s_name] = None\n        finally:\n            for obj in [\"data\", \"atom_array\", \"raw_coords\"]: locals().pop(obj, None)\n            gc.collect(); torch.cuda.empty_cache()\n\n    protenix_preds = {}\n    for tid, info in protenix_queue.items():\n        n_infer, chunks = infer_slots_map[tid], chunk_info.get(tid, [])\n        if len(chunks) == 1:\n            if raw_predictions.get(tid) is not None:\n                protenix_preds[tid] = sanitize_stack(raw_predictions[tid], tid, GLOBAL_SEGMENTS_MAP, confidence=SHORT_PTX_BASE_QUALITY)\n        else:\n            c_res = {s: [] for s in range(n_infer)}\n            for cinfo in chunks:\n                cc = raw_predictions.get(cinfo[\"name\"])\n                if cc is not None:\n                    for s in range(n_infer): c_res[s].append((cc[s] if s < cc.shape[0] else cc[-1], cinfo[\"range\"]))\n            if len(c_res[0]) == len(chunks):\n                stitched = np.stack([stitch_chunk_coords([c for c, _ in c_res[s]], [r for _, r in c_res[s]], len(info[\"sequence\"])) for s in range(n_infer)], axis=0)\n                protenix_preds[tid] = sanitize_stack(stitched, tid, GLOBAL_SEGMENTS_MAP, confidence=LONG_PTX_BASE_QUALITY)\n    return protenix_preds\n\ndef assemble_submission(test_df: pd.DataFrame, plans: dict, protenix_preds: dict, segments_map: dict) -> pd.DataFrame:\n    print(f\"\\n{'=' * 72}\\nPHASE 3: final candidate selection + submission writeout\\n{'=' * 72}\")\n    all_rows = []\n    for r_idx, row in test_df.reset_index(drop=True).iterrows():\n        tid, seq, plan = row[\"target_id\"], row[\"sequence\"], plans[row[\"target_id\"]]\n        pool, tmpls = list(plan[\"template_candidates\"]), list(plan[\"template_candidates\"])\n        \n        if tmpls and len(tmpls) < plan[\"route_info\"][\"template_slots\"]:\n            pool.extend(make_template_aug_candidates(tmpls[0], tid, segments_map, min(AUGMENT_TEMPLATE_VARIANTS, plan[\"route_info\"][\"template_slots\"] - len(tmpls)), (r_idx + 1) * 1009))\n\n        ptx = protenix_preds.get(tid)\n        if ptx is not None and ptx.ndim == 3:\n            # --- PROMPT 3: DYNAMIC QUALITY ---\n            best_tmpl_conf = tmpls[0][\"quality\"] if tmpls else 0.0\n            if best_tmpl_conf >= 0.55: base_q = 0.60 if len(seq) <= MAX_SEQ_LEN else 0.54\n            elif best_tmpl_conf == 0.0: base_q = 0.72 if len(seq) <= MAX_SEQ_LEN else 0.66\n            else: base_q = SHORT_PTX_BASE_QUALITY if len(seq) <= MAX_SEQ_LEN else LONG_PTX_BASE_QUALITY\n\n            for j in range(ptx.shape[0]):\n                pool.append({\"pred_id\": f\"protenix::{j}\", \"source\": \"protenix\", \"quality\": max(0.05, base_q - 0.03 * j), \"coords\": ptx[j]})\n            \n            need_ptx = max(0, plan[\"route_info\"][\"protenix_slots\"] - ptx.shape[0])\n            if need_ptx > 0: pool.extend(make_prediction_aug_candidates(ptx[0], tid, segments_map, min(need_ptx, 3), (r_idx + 1) * 7919, \"protenix_aug\", base_q - 0.01))\n\n        d_idx = 0\n        while len(pool) < max(N_SAMPLE, 7):\n            pool.append(make_de_novo_candidate(seq, tid, segments_map, d_idx, (r_idx + 1) * 1000003 + d_idx * 7919))\n            d_idx += 1\n\n        final_cands = select_final_five_scored(pool, N_SAMPLE)\n        while len(final_cands) < N_SAMPLE:\n            final_cands.append(make_de_novo_candidate(seq, tid, segments_map, len(final_cands), (r_idx + 1) * 1409))\n\n        stacked = np.stack([sanitize_coords(c[\"coords\"], tid, segments_map, c.get(\"quality\", 0.5), source=c.get(\"source\", \"protenix\")) for c in final_cands[:N_SAMPLE]], axis=0)\n        print(f\"  {tid}: \" + \", \".join(f\"{i+1}:{c['source']}({c['quality']:.3f})\" for i, c in enumerate(final_cands[:N_SAMPLE])))\n        all_rows.extend(coords_to_rows(tid, seq, stacked))\n\n    sub = pd.DataFrame(all_rows)\n    cols = [\"ID\", \"resname\", \"resid\"] + [f\"{c}_{i}\" for i in range(1, N_SAMPLE + 1) for c in [\"x\", \"y\", \"z\"]]\n    sub[cols[3:]] = sub[cols[3:]].replace([np.inf, -np.inf], np.nan).fillna(0.0).clip(COORD_CLIP_MIN, COORD_CLIP_MAX)\n    return sub[cols]\n\nGLOBAL_SEGMENTS_MAP = {}\ndef main() -> None:\n    global GLOBAL_SEGMENTS_MAP, USE_PROTENIX\n    test_csv, train_csv, train_lbls, val_csv, val_lbls, output_csv, code_dir, root_dir = resolve_paths()\n    if USE_PROTENIX:\n        os.environ[\"PROTENIX_ROOT_DIR\"] = ensure_required_files(root_dir)\n        if code_dir not in sys.path: sys.path.append(code_dir)\n    seed_everything(SEED)\n    test_df = pd.read_csv(test_csv).head(LOCAL_N_SAMPLES) if not IS_KAGGLE and LOCAL_N_SAMPLES else pd.read_csv(test_csv)\n    c_seqs = pd.concat([pd.read_csv(train_csv), pd.read_csv(val_csv)], ignore_index=True)\n    GLOBAL_SEGMENTS_MAP, _ = build_segments_map(test_df)\n    plans, p_queue = plan_targets(test_df, c_seqs, process_labels(pd.concat([pd.read_csv(train_lbls), pd.read_csv(val_lbls)], ignore_index=True)), GLOBAL_SEGMENTS_MAP)\n    sub = assemble_submission(test_df, plans, run_protenix_phase(p_queue), GLOBAL_SEGMENTS_MAP)\n    Path(output_csv).parent.mkdir(parents=True, exist_ok=True)\n    sub.to_csv(output_csv, index=False)\n    print(f\"\\n✓ Saved submission ({len(sub):,} rows)\")\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}