{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"datasetVersion","sourceId":14604295,"datasetId":9328538,"databundleVersionId":15440074},{"sourceType":"datasetVersion","sourceId":14962460,"datasetId":9577079,"databundleVersionId":15833819},{"sourceType":"datasetVersion","sourceId":14874339,"datasetId":9502242,"databundleVersionId":15736806},{"sourceType":"datasetVersion","sourceId":14962495,"datasetId":9577097,"databundleVersionId":15833858},{"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":"# =========================================================\n# CHUNK 1: installs and imports\n# =========================================================\n\n!pip install --no-index --no-deps /kaggle/input/datasets/kami1976/biopython-cp312/biopython-1.86-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl\n!pip install --no-index --no-deps /kaggle/input/datasets/amirrezaaleyasin/biotite/biotite-1.6.0-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl\n!pip install --no-index --no-deps /kaggle/input/datasets/amirrezaaleyasin/rdkit-2025-9-5/rdkit-2025.9.5-cp312-cp312-manylinux_2_28_x86_64.whl\n\nimport gc\nimport json\nimport os\nimport sys\nimport time\nfrom pathlib import Path\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom Bio.Align import PairwiseAligner\nfrom tqdm import tqdm","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 2: configuration and constants\n# =========================================================\n\nIS_KAGGLE = True\nLOCAL_N_SAMPLES = None\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\nos.environ[\"LAYERNORM_TYPE\"] = \"torch\"\nos.environ.setdefault(\"RNA_MSA_DEPTH_LIMIT\", \"512\")\n\nDATA_BASE              = \"/kaggle/input/competitions/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\", \"128\"))\n\nMIN_SIMILARITY       = float(os.environ.get(\"MIN_SIMILARITY\", \"0.55\"))\nMIN_PERCENT_IDENTITY = float(os.environ.get(\"MIN_PERCENT_IDENTITY\", \"55.0\"))\nTBM_STRONG_SIM       = float(os.environ.get(\"TBM_STRONG_SIM\", \"0.72\"))\nTBM_STRONG_PID       = float(os.environ.get(\"TBM_STRONG_PID\", \"70.0\"))\nTBM_MAX_KEEP         = int(os.environ.get(\"TBM_MAX_KEEP\", \"4\"))\nALWAYS_RUN_PROTENIX  = str(os.environ.get(\"ALWAYS_RUN_PROTENIX\", \"false\")).lower() in {\"1\", \"true\", \"yes\", \"y\"}\n\nUSE_PROTENIX = True\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\"))\nMODEL_N_SAMPLE = int(os.environ.get(\"MODEL_N_SAMPLE\", str(N_SAMPLE)))","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 3: general utility functions\n# =========================================================\n\ndef seed_everything(seed: int) -> None:\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    required = [\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    for p, name in required:\n        if not p.exists():\n            raise FileNotFoundError(f\"Missing {name}: {p}\")\n\n\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 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            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\ndef split_into_chunks(seq_len: int, max_len: int, overlap: int) -> list:\n    if seq_len <= max_len:\n        return [(0, seq_len)]\n    chunks = []\n    step = max_len - overlap\n    pos = 0\n    while pos < seq_len:\n        end = min(pos + max_len, seq_len)\n        chunks.append((pos, end))\n        if end == seq_len:\n            break\n        pos += step\n    return chunks\n\n\ndef kabsch_align(P: np.ndarray, Q: np.ndarray):\n    centroid_P = P.mean(axis=0)\n    centroid_Q = Q.mean(axis=0)\n    Pc = P - centroid_P\n    Qc = Q - centroid_Q\n    H = Pc.T @ Qc\n    U, _, Vt = np.linalg.svd(H)\n    d = np.linalg.det(Vt.T @ U.T)\n    S = np.eye(3)\n    if d < 0:\n        S[2, 2] = -1\n    R = Vt.T @ S @ U.T\n    t = centroid_Q - R @ centroid_P\n    return R, t\n\n\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:\n            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\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\n        ov_start = cur_start\n        ov_end = min(prev_end, cur_end)\n        ov_len = ov_end - ov_start\n\n        if ov_len < 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\n        valid = ~(np.isnan(prev_ov).any(axis=1) | np.isnan(cur_ov).any(axis=1))\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        transformed = (chunk_coords_list[i] @ R.T) + t\n        aligned.append(transformed)\n\n    full = np.zeros((seq_len, 3), dtype=np.float64)\n    weights = np.zeros(seq_len, dtype=np.float64)\n\n    for i, ((s, e), coords) in enumerate(zip(chunk_ranges, aligned)):\n        chunk_len = coords.shape[0]\n        actual_end = min(s + chunk_len, seq_len)\n        used_len = actual_end - s\n\n        w = np.ones(used_len, dtype=np.float64)\n\n        if i > 0:\n            ov_start = s\n            ov_end = min(chunk_ranges[i - 1][1], e)\n            ramp_len = ov_end - ov_start\n            if ramp_len > 0:\n                w[:ramp_len] = np.linspace(0.0, 1.0, ramp_len)\n\n        if i < len(chunk_ranges) - 1:\n            next_s = chunk_ranges[i + 1][0]\n            ramp_start = next_s - s\n            ramp_len = actual_end - next_s\n            if ramp_len > 0 and ramp_start < used_len:\n                w[ramp_start:used_len] = np.linspace(1.0, 0.0, ramp_len)\n\n        full[s:actual_end] += coords[:used_len] * w[:, None]\n        weights[s:actual_end] += w\n\n    mask = weights > 0\n    full[mask] /= weights[mask, None]\n\n    return full","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 4: RNA geometry / augmentation / fallback functions\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\n    for _ in range(passes):\n        for s, e in segments:\n            C = X[s:e]\n            L = e - s\n            if L < 3:\n                continue\n\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\n            C[1:] += adj\n\n            d2 = C[2:] - C[:-2]\n            d2n = np.linalg.norm(d2, axis=1) + 1e-6\n            adj2 = d2 * ((10.2 - d2n) / d2n)[:, None] * (0.10 * strength)\n            C[:-2] -= adj2\n            C[2:] += adj2\n\n            C[1:-1] += (0.06 * strength) * (0.5 * (C[:-2] + C[2:]) - C[1:-1])\n\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]\n                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\n            X[s:e] = C\n\n    return X\n\n\ndef _rotmat(axis, ang):\n    a = np.asarray(axis, float)\n    a /= np.linalg.norm(a) + 1e-12\n    x, y, z = a\n    c, s = np.cos(ang), np.sin(ang)\n    CC = 1 - c\n    return np.array([\n        [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\n\ndef apply_hinge(coords, seg, rng, deg=22):\n    s, e = seg\n    L = e - s\n    if L < 30:\n        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()\n    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    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)\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    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:\n            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 generate_rna_structure(sequence: str, seed=None) -> np.ndarray:\n    if seed is not None:\n        np.random.seed(seed)\n    n = len(sequence)\n    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","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 5: TBM parsing and alignment functions\n# =========================================================\n\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\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\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\n    for (qs, qe), (ts, te) in zip(q_segs, t_segs):\n        while qi < qs:\n            aq.append(query_seq[qi])\n            at.append(\"-\")\n            qi += 1\n        while ti < ts:\n            aq.append(\"-\")\n            at.append(template_seq[ti])\n            ti += 1\n        for qp, tp in zip(range(qs, qe), range(ts, te)):\n            aq.append(query_seq[qp])\n            at.append(template_seq[tp])\n        qi, ti = qe, te\n\n    while qi < len(query_seq):\n        aq.append(query_seq[qi])\n        at.append(\"-\")\n        qi += 1\n    while ti < len(template_seq):\n        aq.append(\"-\")\n        at.append(template_seq[ti])\n        ti += 1\n\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    qlen = len(query_seq)\n\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\n        tlen = len(tseq)\n        if max(qlen, tlen) == 0:\n            continue\n        if abs(tlen - qlen) / max(tlen, qlen) > 0.35:\n            continue\n\n        aln = next(iter(_aligner.align(query_seq, tseq)))\n        aq, at = _build_aligned_strings(query_seq, tseq, aln)\n\n        aligned_pairs = 0\n        identical = 0\n        gap_q = 0\n        gap_t = 0\n\n        for a, b in zip(aq, at):\n            if a != '-' and b != '-':\n                aligned_pairs += 1\n                identical += int(a == b)\n            elif a == '-':\n                gap_q += 1\n            elif b == '-':\n                gap_t += 1\n\n        coverage = aligned_pairs / max(1, qlen)\n        pct_id = 100.0 * identical / max(1, qlen)\n        norm_s = aln.score / (2.0 * max(1, min(qlen, tlen)))\n        len_ratio = min(qlen, tlen) / max(qlen, tlen)\n        gap_penalty = (gap_q + gap_t) / max(1, qlen + tlen)\n\n        hybrid_score = (\n            0.45 * norm_s\n            + 0.25 * (pct_id / 100.0)\n            + 0.20 * coverage\n            + 0.10 * len_ratio\n            - 0.15 * gap_penalty\n        )\n\n        results.append({\n            \"target_id\": tid,\n            \"sequence\": tseq,\n            \"norm_score\": float(norm_s),\n            \"pct_identity\": float(pct_id),\n            \"coverage\": float(coverage),\n            \"len_ratio\": float(len_ratio),\n            \"hybrid_score\": float(hybrid_score),\n            \"coords\": train_coords_dict[tid],\n            \"aligned_query\": aq,\n            \"aligned_template\": at,\n        })\n\n    results.sort(key=lambda x: (x[\"hybrid_score\"], x[\"pct_identity\"], x[\"coverage\"]), 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\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\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\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\n    return np.nan_to_num(new_coords)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 6: TBM phase\n# =========================================================\n\ndef tbm_phase(test_df, train_seqs_df, train_coords_dict, segments_map):\n    print(\"\\n\" + \"=\"*60)\n    print(\"PHASE 1: Template-Based Modeling\")\n    print(\n        f\"  MIN_SIMILARITY = {MIN_SIMILARITY} | MIN_PCT_IDENTITY = {MIN_PERCENT_IDENTITY} \"\n        f\"| TBM_STRONG_SIM = {TBM_STRONG_SIM} | TBM_STRONG_PID = {TBM_STRONG_PID}\"\n    )\n    print(\"=\"*60)\n\n    t0 = time.time()\n    template_predictions = {}\n    protenix_queue = {}\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=24)\n\n        preds = []\n        kept_templates = []\n\n        for cand in similar:\n            sim = cand[\"norm_score\"]\n            pct_id = cand[\"pct_identity\"]\n            cov = cand[\"coverage\"]\n\n            if sim < MIN_SIMILARITY or pct_id < MIN_PERCENT_IDENTITY or cov < 0.72:\n                continue\n            if len(preds) >= TBM_MAX_KEEP:\n                break\n\n            tmpl_id = cand[\"target_id\"]\n            tmpl_seq = cand[\"sequence\"]\n            tmpl_coords = cand[\"coords\"]\n\n            rng = np.random.default_rng(abs(hash((tid, tmpl_id, len(preds), SEED))) % (2**32))\n            adapted = adapt_template_to_query(seq, tmpl_seq, tmpl_coords)\n\n            \n            slot = len(preds)\n            if slot == 0:\n                X = adapted\n            else:\n                X = adapted.copy()\n                X = X + rng.normal(0, 0.15 + 0.10 * slot, adapted.shape)\n\n            preds.append(X)\n            kept_templates.append((tmpl_id, sim, pct_id, cov, cand[\"hybrid_score\"]))\n\n        template_predictions[tid] = preds\n\n        if ALWAYS_RUN_PROTENIX:\n            n_needed = N_SAMPLE\n        else:\n            n_needed = N_SAMPLE - len(preds)\n\n        protenix_queue[tid] = (n_needed, seq)\n\n        if kept_templates:\n            desc = \", \".join(\n                [f\"{x[0]} sim={x[1]:.3f} pid={x[2]:.1f}% cov={x[3]:.2f}\" for x in kept_templates]\n            )\n            print(f\"  {tid} ({len(seq)} nt): kept {len(preds)} TBM anchor(s) [{desc}] + Protenix {n_needed}\")\n        else:\n            print(f\"  {tid} ({len(seq)} nt): no strong TBM anchor kept + Protenix {n_needed}\")\n\n    elapsed = time.time() - t0\n    print(f\"\\nPhase 1 done in {elapsed:.1f}s\")\n    print(f\"  Targets with TBM anchors : {sum(len(v) > 0 for v in template_predictions.values())}\")\n    print(f\"  Targets sent to Protenix : {len(protenix_queue)}\")\n\n    return template_predictions, protenix_queue","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 7: Protenix config and helper functions\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\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 extract_protenix_c1_coords(prediction, feat, chunk_seq_len, raw_coords):\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        m11 = (feat[\"atom_to_tokatom_idx\"] == 11).to(raw_coords.device)\n        m12 = (feat[\"atom_to_tokatom_idx\"] == 12).to(raw_coords.device)\n        c11, c12 = m11.sum(), m12.sum()\n        mask = m11 if abs(c11 - chunk_seq_len) < abs(c12 - chunk_seq_len) else m12\n    else:\n        mask = torch.zeros(raw_coords.shape[1], dtype=torch.bool, device=raw_coords.device)\n\n    coords = raw_coords[:, mask, :].detach().cpu().numpy()\n\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(\"    WARNING: collapsed coordinates detected\")\n            return None\n\n    if coords.shape[1] != chunk_seq_len:\n        if coords.shape[1] == 1 and chunk_seq_len > 1:\n            return None\n        padded = np.zeros((coords.shape[0], chunk_seq_len, 3), dtype=np.float32)\n        ml = min(coords.shape[1], chunk_seq_len)\n        padded[:, :ml, :] = coords[:, :ml, :]\n        coords = padded\n\n    return coords","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 8: Protenix inference phase\n# =========================================================\n\ndef protenix_phase(protenix_queue, work_dir=\"/kaggle/working\"):\n    protenix_preds = {}\n\n    if not protenix_queue or not USE_PROTENIX:\n        return protenix_preds\n\n    print(\"\\n\" + \"=\"*60)\n    print(f\"PHASE 2: Protenix for {len(protenix_queue)} targets\")\n    print(\"=\"*60)\n\n    work_dir = Path(work_dir)\n    work_dir.mkdir(parents=True, exist_ok=True)\n\n    tasks = []\n    chunk_info = {}\n\n    for target_id, (n_needed, full_seq) in protenix_queue.items():\n        seq_len = len(full_seq)\n\n        if seq_len <= MAX_SEQ_LEN:\n            tasks.append({\"target_id\": target_id, \"sequence\": full_seq})\n            chunk_info[target_id] = [{\"name\": target_id, \"range\": (0, seq_len)}]\n            print(f\"  {target_id} ({seq_len} nt): single pass queued\")\n        else:\n            chunks = split_into_chunks(seq_len, MAX_SEQ_LEN, CHUNK_OVERLAP)\n            print(f\"  {target_id} ({seq_len} nt): {len(chunks)} chunks queued {[(s, e) for s, e in chunks]}\")\n\n            chunk_info[target_id] = []\n            for ci, (cs, ce) in enumerate(chunks):\n                chunk_name = f\"{target_id}_chunk{ci}\"\n                sub_seq = full_seq[cs:ce]\n                tasks.append({\"target_id\": chunk_name, \"sequence\": sub_seq})\n                chunk_info[target_id].append({\"name\": chunk_name, \"range\": (cs, ce)})\n\n    tasks_df = pd.DataFrame(tasks)\n    input_json_path = str(work_dir / \"protenix_queue_input.json\")\n    build_input_json(tasks_df, input_json_path)\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 = 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    raw_predictions = {}\n\n    for i in tqdm(range(len(dataset)), desc=\"Protenix Inference\"):\n        data, atom_array, err = dataset[i]\n        sample_name = data.get(\"sample_name\", f\"sample_{i}\")\n\n        if err:\n            print(f\"  {sample_name} data error: {err}\")\n            raw_predictions[sample_name] = None\n            del data, atom_array, err\n            gc.collect()\n            torch.cuda.empty_cache()\n            gc.collect()\n            continue\n\n        target_id = sample_name.split(\"_chunk\")[0] if \"_chunk\" in sample_name else sample_name\n        n_needed = protenix_queue.get(target_id, (N_SAMPLE, \"\"))[0]\n        sub_seq_len = data[\"N_token\"].item()\n\n        try:\n            new_cfg = update_inference_configs(configs, sub_seq_len)\n            new_cfg.sample_diffusion.N_sample = n_needed\n            runner.update_model_configs(new_cfg)\n\n            pred = runner.predict(data)\n            raw_coords = pred[\"coordinate\"]\n\n            coords = extract_protenix_c1_coords(\n                pred,\n                data[\"input_feature_dict\"],\n                sub_seq_len,\n                raw_coords\n            )\n            raw_predictions[sample_name] = coords\n\n        except Exception as exc:\n            print(f\"  {sample_name} inference failed: {exc}\")\n            import traceback\n            traceback.print_exc()\n            raw_predictions[sample_name] = None\n\n        finally:\n            try:\n                del pred, data, atom_array, raw_coords\n            except:\n                pass\n            gc.collect()\n            torch.cuda.empty_cache()\n            gc.collect()\n\n    for target_id, (n_needed, full_seq) in protenix_queue.items():\n        seq_len = len(full_seq)\n        chunks = chunk_info.get(target_id, [])\n\n        if not chunks:\n            continue\n\n        if len(chunks) == 1:\n            coords = raw_predictions.get(target_id)\n            protenix_preds[target_id] = coords\n            if coords is not None:\n                print(f\"  {target_id}: {coords.shape[0]} predictions generated\")\n            else:\n                print(f\"  {target_id}: FAILED\")\n        else:\n            chunk_results_per_sample = {s: [] for s in range(n_needed)}\n            all_ok = True\n\n            for ci, cinfo in enumerate(chunks):\n                cname = cinfo[\"name\"]\n                crange = cinfo[\"range\"]\n                ccoords = raw_predictions.get(cname)\n\n                if ccoords is None:\n                    all_ok = False\n                    break\n\n                for s_idx in range(n_needed):\n                    if s_idx < ccoords.shape[0]:\n                        chunk_results_per_sample[s_idx].append((ccoords[s_idx], crange))\n                    else:\n                        chunk_results_per_sample[s_idx].append((ccoords[-1], crange))\n\n            if not all_ok:\n                print(f\"  {target_id}: chunked inference incomplete, using fallback\")\n                protenix_preds[target_id] = None\n                continue\n\n            stitched_samples = []\n            for s_idx in range(n_needed):\n                items = chunk_results_per_sample[s_idx]\n                coords_list = [c for c, _ in items]\n                ranges_list = [r for _, r in items]\n                full_coords = stitch_chunk_coords(coords_list, ranges_list, seq_len)\n                stitched_samples.append(full_coords)\n\n            result = np.stack(stitched_samples, axis=0)\n            protenix_preds[target_id] = result\n            print(f\"  {target_id}: {result.shape[0]} stitched predictions generated\")\n\n    return protenix_preds","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 9: combine predictions and build submission\n# =========================================================\n\ndef combine_predictions(test_df, template_preds, protenix_preds, segments_map):\n    all_rows = []\n\n    for _, row in test_df.iterrows():\n        tid = row[\"target_id\"]\n        seq = row[\"sequence\"]\n\n        tbm_list = list(template_preds.get(tid, []))\n        combined = []\n\n        if tbm_list:\n            combined.append(tbm_list[0])\n\n        ptx = protenix_preds.get(tid)\n        if ptx is not None and getattr(ptx, \"ndim\", 0) == 3:\n            max_ptx_to_add = N_SAMPLE - len(combined) - (1 if len(tbm_list) > 1 else 0)\n            max_ptx_to_add = max(0, min(max_ptx_to_add, ptx.shape[0]))\n            for j in range(max_ptx_to_add):\n                combined.append((ptx[j])\n\n        if len(tbm_list) > 1 and len(combined) < N_SAMPLE:\n            combined.append(tbm_list[1])\n\n        n_denovo = 0\n        while len(combined) < N_SAMPLE:\n            seed_val = row.name * 1000000 + len(combined) * 1000\n            dn = generate_rna_structure(seq, seed=seed_val)\n            combined.append(dn)\n            n_denovo += 1\n\n        if n_denovo:\n            print(f\"  {tid}: {n_denovo} slot(s) filled with de-novo fallback\")\n\n        stacked = np.stack(combined[:N_SAMPLE], axis=0)\n        all_rows.extend(coords_to_rows(tid, seq, stacked))\n\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    return sub[cols]\n\n\ndef save_submission(sub_df, output_csv):\n    sub_df.to_csv(output_csv, index=False)\n    print(f\"\\n✓ Saved submission to {output_csv} ({len(sub_df):,} rows)\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================\n# CHUNK 10: load data and main runner\n# =========================================================\n\ndef load_competition_data(test_csv):\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    print(f\"Test targets : {len(test_df)}\" + (\" (LOCAL MODE)\" if not IS_KAGGLE else \"\"))\n\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\n    print(f\"Template pool: {len(combined_seqs)} sequences, {len(train_coords)} structures\")\n    return test_df, combined_seqs, train_coords\n\n\ndef prepare_protenix_env(code_dir, root_dir):\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    os.environ[\"PROTENIX_ROOT_DIR\"] = root_dir\n    sys.path.append(code_dir)\n    ensure_required_files(root_dir)\n\n\ndef main():\n    test_csv, output_csv, code_dir, root_dir = resolve_paths()\n\n    prepare_protenix_env(code_dir, root_dir)\n    seed_everything(SEED)\n\n    test_df, combined_seqs, train_coords = load_competition_data(test_csv)\n    segments_map, _ = build_segments_map(test_df)\n\n    template_preds, protenix_queue = tbm_phase(\n        test_df=test_df,\n        train_seqs_df=combined_seqs,\n        train_coords_dict=train_coords,\n        segments_map=segments_map\n    )\n\n    if protenix_queue and USE_PROTENIX:\n        protenix_preds = protenix_phase(protenix_queue, work_dir=\"/kaggle/working\")\n    else:\n        protenix_preds = {}\n        if protenix_queue:\n            print(f\"\\nPHASE 2 skipped (USE_PROTENIX=False). De-novo fallback will cover {len(protenix_queue)} targets.\")\n\n    print(\"\\n\" + \"=\"*60)\n    print(\"PHASE 3: Combine TBM + Protenix + de-novo fallback\")\n    print(\"=\"*60)\n\n    sub_df = combine_predictions(\n        test_df=test_df,\n        template_preds=template_preds,\n        protenix_preds=protenix_preds,\n        segments_map=segments_map\n    )\n\n    save_submission(sub_df, output_csv)\n\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}