{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":118765,"databundleVersionId":15231210},{"sourceType":"datasetVersion","sourceId":14970008,"datasetId":9582207,"databundleVersionId":15842078},{"sourceType":"datasetVersion","sourceId":14874339,"datasetId":9502242,"databundleVersionId":15736806}],"dockerImageVersionId":31287,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"WHEEL_DIR = \"/kaggle/input/datasets/esabellechen/biopython-wheel\"\n\n!pip install --no-index --find-links=$WHEEL_DIR \\\n    numpy scipy networkx packaging pillow msgpack \\\n    certifi charset_normalizer idna urllib3 requests \\\n    biopython rdkit biotraj biotite","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:31:05.229874Z","iopub.execute_input":"2026-03-25T09:31:05.230168Z","iopub.status.idle":"2026-03-25T09:31:15.973559Z","shell.execute_reply.started":"2026-03-25T09:31:05.230115Z","shell.execute_reply":"2026-03-25T09:31:15.972841Z"},"jupyter":{"source_hidden":true}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# stitch_chunks_graph","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:31:15.976303Z","iopub.execute_input":"2026-03-25T09:31:15.9766Z","iopub.status.idle":"2026-03-25T09:31:15.980449Z","shell.execute_reply.started":"2026-03-25T09:31:15.976569Z","shell.execute_reply":"2026-03-25T09:31:15.979766Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nimport sys\nimport pandas as pd\n\n# ── Local vs Kaggle mode ─────────────────────────────────────────────────────\n# On Kaggle competition rerun, KAGGLE_IS_COMPETITION_RERUN is set to a truthy value.\n# When running locally we do NOT exit — instead we cap the test set to a small\n# number of samples so the notebook finishes quickly.\n\nIS_KAGGLE = bool(os.environ.get(\"KAGGLE_IS_COMPETITION_RERUN\", \"\"))\n\n# How many test samples to use when running locally\nLOCAL_N_SAMPLES = 2\nSINGLE_TEST = '9ZCC' #9MME: 4640 9ZCC: 1460\nDO_SINGLE_TEST = False\n\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 \"\n          f\"will be processed to save time.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:31:15.981597Z","iopub.execute_input":"2026-03-25T09:31:15.981819Z","iopub.status.idle":"2026-03-25T09:31:16.26752Z","shell.execute_reply.started":"2026-03-25T09:31:15.981798Z","shell.execute_reply":"2026-03-25T09:31:16.266785Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\nimport json\nimport os\nimport time\nimport random\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","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:31:16.268512Z","iopub.execute_input":"2026-03-25T09:31:16.268986Z","iopub.status.idle":"2026-03-25T09:31:20.589136Z","shell.execute_reply.started":"2026-03-25T09:31:16.268961Z","shell.execute_reply":"2026-03-25T09:31:20.588435Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\n# ─────────────── Paths & Constants ───────────────────────────────────────────\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      = 10\nFINAL_N_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\n# TBM quality thresholds — sequences below these get routed to Protenix\nMIN_SIMILARITY       = float(os.environ.get(\"MIN_SIMILARITY\",       \"0.0\"))\nMIN_PERCENT_IDENTITY = float(os.environ.get(\"MIN_PERCENT_IDENTITY\", \"50.0\"))\n\n# Set False to skip Protenix and use de-novo fallback instead\nUSE_PROTENIX = True\n\n\ndef parse_bool(value: str, default: bool = False) -> str:\n    v = str(value).strip().lower()\n    if v in {\"1\", \"true\", \"t\", \"yes\", \"y\", \"on\"}:\n        return \"true\"\n    if v in {\"0\", \"false\", \"f\", \"no\", \"n\", \"off\"}:\n        return \"false\"\n    return \"true\" if default else \"false\"\n\n\nUSE_MSA      = parse_bool(os.environ.get(\"USE_MSA\",      \"false\"))\nUSE_TEMPLATE = parse_bool(os.environ.get(\"USE_TEMPLATE\", \"false\"))\nUSE_RNA_MSA  = parse_bool(os.environ.get(\"USE_RNA_MSA\",  \"true\"))\n\nMODEL_N_SAMPLE = int(os.environ.get(\"MODEL_N_SAMPLE\", str(N_SAMPLE)))\n\n\n# ─────────────── General Utilities ───────────────────────────────────────────\ndef seed_everything(seed: int) -> None:\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    os.environ[\"CUBLAS_WORKSPACE_CONFIG\"] = \":4096:8\"\n    random.seed(SEED)\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\n\ndef resolve_paths():\n    test_csv   = os.environ.get(\"TEST_CSV\",           DEFAULT_TEST_CSV)\n    output_csv = os.environ.get(\"SUBMISSION_CSV\",     DEFAULT_OUTPUT)\n    code_dir   = os.environ.get(\"PROTENIX_CODE_DIR\",  DEFAULT_CODE_DIR)\n    root_dir   = os.environ.get(\"PROTENIX_ROOT_DIR\",  DEFAULT_ROOT_DIR)\n    return test_csv, output_csv, code_dir, root_dir\n\n\ndef ensure_required_files(root_dir: str) -> None:\n    for p, name in [\n        (Path(root_dir) / \"checkpoint\" / f\"{MODEL_NAME}.pt\",          \"checkpoint\"),\n        (Path(root_dir) / \"common\" / \"components.cif\",                \"CCD file\"),\n        (Path(root_dir) / \"common\" / \"components.cif.rdkit_mol.pkl\",  \"CCD cache\"),\n    ]:\n        if not p.exists():\n            raise FileNotFoundError(f\"Missing {name}: {p}\")\n\n\n# ─────────────── Protenix Input / Config Helpers ─────────────────────────────\ndef build_input_json(df: pd.DataFrame, json_path: str) -> None:\n    data = [\n        {\n            \"name\": row[\"target_id\"],\n            \"covalent_bonds\": [],\n            \"sequences\": [{\"rnaSequence\": {\"sequence\": row[\"sequence\"], \"count\": 1}}],\n        }\n        for _, row in df.iterrows()\n    ]\n    with open(json_path, \"w\", encoding=\"utf-8\") as f:\n        json.dump(data, f)\n\n\ndef build_configs(input_json_path: str, dump_dir: str, model_name: str):\n    from configs.configs_base import configs as configs_base\n    from configs.configs_data import data_configs\n    from configs.configs_inference import inference_configs\n    from configs.configs_model_type import model_configs\n    from protenix.config.config import parse_configs\n\n    base = {**configs_base, **{\"data\": data_configs}, **inference_configs}\n\n    def deep_update(t, p):\n        for k, v in p.items():\n            if isinstance(v, dict) and k in t and isinstance(t[k], dict):\n                deep_update(t[k], v)\n            else:\n                t[k] = v\n\n    deep_update(base, model_configs[model_name])\n    arg_str = \" \".join([\n        f\"--model_name {model_name}\",\n        f\"--input_json_path {input_json_path}\",\n        f\"--dump_dir {dump_dir}\",\n        f\"--use_msa {USE_MSA}\",\n        f\"--use_template {USE_TEMPLATE}\",\n        f\"--use_rna_msa {USE_RNA_MSA}\",\n        f\"--sample_diffusion.N_sample {MODEL_N_SAMPLE}\",\n        f\"--seeds {SEED}\",\n    ])\n    return parse_configs(configs=base, arg_str=arg_str, fill_required_with_null=True)\n\n\n\ndef coords_to_rows(target_id: str, seq: str, coords: np.ndarray) -> list:\n    \"\"\"coords shape: (N_SAMPLE, seq_len, 3)\"\"\"\n    rows = []\n    for i in range(len(seq)):\n        row = {\"ID\": f\"{target_id}_{i + 1}\", \"resname\": seq[i], \"resid\": i + 1}\n        for s in range(N_SAMPLE):\n            if s < coords.shape[0] and i < coords.shape[1]:\n                x, y, z = coords[s, i]\n            else:\n                x, y, z = 0.0, 0.0, 0.0\n            row[f\"x_{s + 1}\"] = float(x)\n            row[f\"y_{s + 1}\"] = float(y)\n            row[f\"z_{s + 1}\"] = float(z)\n        rows.append(row)\n    return rows\n\n\n\n# ─────────────── TBM Core Functions ──────────────────────────────────────────\ndef _make_aligner() -> PairwiseAligner:\n    al = PairwiseAligner()\n    al.mode                           = \"global\"\n    al.match_score                    = 2\n    al.mismatch_score                 = -1.5\n    al.open_gap_score                 = -8\n    al.extend_gap_score               = -0.4\n    al.query_left_open_gap_score      = -8\n    al.query_left_extend_gap_score    = -0.4\n    al.query_right_open_gap_score     = -8\n    al.query_right_extend_gap_score   = -0.4\n    al.target_left_open_gap_score     = -8\n    al.target_left_extend_gap_score   = -0.4\n    al.target_right_open_gap_score    = -8\n    al.target_right_extend_gap_score  = -0.4\n    return al\n\n\n_aligner = _make_aligner()\n\n\ndef parse_stoichiometry(stoich: str) -> list:\n    if pd.isna(stoich) or str(stoich).strip() == \"\":\n        return []\n    return [(ch.strip(), int(cnt)) for part in str(stoich).split(\";\")\n            for ch, cnt in [part.split(\":\")]]\n\n\ndef parse_fasta(fasta_content: str) -> dict:\n    out, cur, parts = {}, None, []\n    for line in str(fasta_content).splitlines():\n        line = line.strip()\n        if not line:\n            continue\n        if line.startswith(\">\"):\n            if cur is not None:\n                out[cur] = \"\".join(parts)\n            cur = line[1:].split()[0]\n            parts = []\n        else:\n            parts.append(line.replace(\" \", \"\"))\n    if cur is not None:\n        out[cur] = \"\".join(parts)\n    return out\n\n\ndef get_chain_segments(row) -> list:\n    seq    = row[\"sequence\"]\n    stoich = row.get(\"stoichiometry\", \"\")\n    all_sq = row.get(\"all_sequences\", \"\")\n    if (pd.isna(stoich) or pd.isna(all_sq)\n            or str(stoich).strip() == \"\" or str(all_sq).strip() == \"\"):\n        return [(0, len(seq))]\n    try:\n        chain_dict = parse_fasta(all_sq)\n        order = parse_stoichiometry(stoich)\n        segs, pos = [], 0\n        for ch, cnt in order:\n            base = chain_dict.get(ch)\n            if base is None:\n                return [(0, len(seq))]\n            for _ in range(cnt):\n                segs.append((pos, pos + len(base)))\n                pos += len(base)\n        return segs if pos == len(seq) else [(0, len(seq))]\n    except Exception:\n        return [(0, len(seq))]\n\n\ndef build_segments_map(df: pd.DataFrame) -> tuple:\n    seg_map, stoich_map = {}, {}\n    for _, r in df.iterrows():\n        tid               = r[\"target_id\"]\n        seg_map[tid]      = get_chain_segments(r)\n        raw_s             = r.get(\"stoichiometry\", \"\")\n        stoich_map[tid]   = \"\" if pd.isna(raw_s) else str(raw_s)\n    return seg_map, stoich_map\n\n\ndef process_labels(labels_df: pd.DataFrame) -> dict:\n    coords = {}\n    prefixes = labels_df[\"ID\"].str.rsplit(\"_\", n=1).str[0]\n    for prefix, grp in labels_df.groupby(prefixes):\n        coords[prefix] = grp.sort_values(\"resid\")[[\"x_1\", \"y_1\", \"z_1\"]].values\n    return coords\n\n\n\ndef find_similar_sequences_detailed(query_seq, train_seqs_df, train_coords_dict, top_n=5):\n    results = []\n    for _, row in train_seqs_df.iterrows():\n        tid, tseq = row[\"target_id\"], row[\"sequence\"]\n        if tid not in train_coords_dict:\n            continue\n        if abs(len(tseq) - len(query_seq)) / max(len(tseq), len(query_seq)) > 0.3:\n            continue\n        aln       = next(iter(_aligner.align(query_seq, tseq)))\n        norm_s    = aln.score / (2 * min(len(query_seq), len(tseq)))\n        identical = sum(\n            1 for (qs, qe), (ts, te) in zip(*aln.aligned)\n            for qp, tp in zip(range(qs, qe), range(ts, te))\n            if query_seq[qp] == tseq[tp]\n        )\n        pct_id = 100 * identical / len(query_seq)\n        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) -> np.ndarray:\n    aln        = next(iter(_aligner.align(query_seq, template_seq)))\n    new_coords = np.full((len(query_seq), 3), np.nan)\n    for (qs, qe), (ts, te) in zip(*aln.aligned):\n        chunk = template_coords[ts:te]\n        if len(chunk) == (qe - qs):\n            new_coords[qs:qe] = chunk\n    for i in range(len(new_coords)):\n        if np.isnan(new_coords[i, 0]):\n            pv = next((j for j in range(i - 1, -1, -1) if not np.isnan(new_coords[j, 0])), -1)\n            nv = next((j for j in range(i + 1, len(new_coords)) if not np.isnan(new_coords[j, 0])), -1)\n            if pv >= 0 and nv >= 0:\n                w = (i - pv) / (nv - pv)\n                new_coords[i] = (1 - w) * new_coords[pv] + w * new_coords[nv]\n            elif pv >= 0:\n                new_coords[i] = new_coords[pv] + [3, 0, 0]\n            elif nv >= 0:\n                new_coords[i] = new_coords[nv] + [3, 0, 0]\n            else:\n                new_coords[i] = [i * 3, 0, 0]\n    return np.nan_to_num(new_coords)\n\n\ndef adaptive_rna_constraints(coords, target_id, segments_map, confidence=1.0, passes=2) -> np.ndarray:\n    X        = coords.copy()\n    segments = segments_map.get(target_id, [(0, len(X))])\n    strength = max(0.75 * (1.0 - min(confidence, 0.97)), 0.02)\n    for _ in range(passes):\n        for s, e in segments:\n            C = X[s:e]; L = e - s\n            if L < 3:\n                continue\n            # bond i–i+1  ~5.95 Å\n            d    = C[1:] - C[:-1]; dist = np.linalg.norm(d, axis=1) + 1e-6\n            adj  = d * ((5.95 - dist) / dist)[:, None] * (0.22 * strength)\n            C[:-1] -= adj; C[1:] += adj\n            # soft i–i+2  ~10.2 Å\n            d2   = C[2:] - C[:-2]; d2n = np.linalg.norm(d2, axis=1) + 1e-6\n            adj2 = d2 * ((10.2 - d2n) / d2n)[:, None] * (0.10 * strength)\n            C[:-2] -= adj2; C[2:] += adj2\n            # Laplacian smoothing\n            C[1:-1] += (0.06 * strength) * (0.5 * (C[:-2] + C[2:]) - C[1:-1])\n            # self-avoidance\n            if L >= 25:\n                idx  = np.linspace(0, L - 1, min(L, 160)).astype(int) if L > 220 else np.arange(L)\n                P    = C[idx]; diff = P[:, None, :] - P[None, :, :]\n                dm   = np.linalg.norm(diff, axis=2) + 1e-6\n                sep  = np.abs(idx[:, None] - idx[None, :])\n                mask = (sep > 2) & (dm < 3.2)\n                if np.any(mask):\n                    vec = (diff * ((3.2 - dm) / dm)[:, :, None] * mask[:, :, None]).sum(axis=1)\n                    C[idx] += (0.015 * strength) * vec\n            X[s:e] = C\n    return X\n\n\ndef _rotmat(axis, ang):\n    a = np.asarray(axis, float); a /= np.linalg.norm(a) + 1e-12\n    x, y, z = a; c, s = np.cos(ang), np.sin(ang); CC = 1 - c\n    return np.array([[c+x*x*CC, x*y*CC-z*s, x*z*CC+y*s],\n                     [y*x*CC+z*s, c+y*y*CC, y*z*CC-x*s],\n                     [z*x*CC-y*s, z*y*CC+x*s, c+z*z*CC]])\n\n\ndef apply_hinge(coords, seg, rng, deg=22):\n    s, e = seg; L = e - s\n    if L < 30: return coords\n    pivot = s + int(rng.integers(10, L - 10))\n    R = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-deg, deg))))\n    X = coords.copy(); p0 = X[pivot].copy()\n    X[pivot+1:e] = (X[pivot+1:e] - p0) @ R.T + p0\n    return X\n\n\ndef jitter_chains(coords, segs, rng, deg=12, trans=1.5):\n    X = coords.copy(); gc_ = X.mean(0, keepdims=True)\n    for s, e in segs:\n        R     = _rotmat(rng.normal(size=3), np.deg2rad(float(rng.uniform(-deg, deg))))\n        shift = rng.normal(size=3); shift = shift / (np.linalg.norm(shift) + 1e-12) * float(rng.uniform(0, trans))\n        c     = X[s:e].mean(0, keepdims=True)\n        X[s:e] = (X[s:e] - c) @ R.T + c + shift\n    X -= X.mean(0, keepdims=True) - gc_\n    return X\n\n\ndef smooth_wiggle(coords, segs, rng, amp=0.8):\n    X = coords.copy()\n    for s, e in segs:\n        L = e - s\n        if L < 20: continue\n        ctrl = np.linspace(0, L - 1, 6); disp = rng.normal(0, amp, (6, 3)); t = np.arange(L)\n        X[s:e] += np.vstack([np.interp(t, ctrl, disp[:, k]) for k in range(3)]).T\n    return X\n\n\ndef generate_rna_structure(sequence: str, seed=None) -> np.ndarray:\n    \"\"\"Idealized A-form RNA helix — last-resort de-novo fallback.\"\"\"\n    if seed is not None:\n        np.random.seed(seed)\n    n = len(sequence); coords = np.zeros((n, 3))\n    for i in range(n):\n        ang = i * 0.6\n        coords[i] = [10.0 * np.cos(ang), 10.0 * np.sin(ang), i * 2.5]\n    return coords\n\n\n\n# ─────────────── TBM Phase ───────────────────────────────────────────────────\ndef tbm_phase(test_df, train_seqs_df, train_coords_dict, segments_map):\n    \"\"\"\n    Phase 1 — Template-Based Modeling.\n\n    Returns\n    -------\n    template_predictions : {target_id: [np.ndarray(seq_len, 3), ...]}\n        0 to N_SAMPLE predictions per target, from real templates.\n    protenix_queue : {target_id: (n_needed, full_sequence)}\n        Targets that still need more predictions.\n    \"\"\"\n    print(f\"\\n{'='*60}\")\n    print(f\"PHASE 1: Template-Based Modeling\")\n    print(f\"  MIN_SIMILARITY = {MIN_SIMILARITY}  |  MIN_PCT_IDENTITY = {MIN_PERCENT_IDENTITY}\")\n    print(f\"{'='*60}\")\n    t0 = time.time()\n\n    template_predictions: dict = {}\n    protenix_queue:       dict = {}\n\n    for _, row in test_df.iterrows():\n        tid = row[\"target_id\"]\n        seq = row[\"sequence\"]\n        segs = segments_map.get(tid, [(0, len(seq))])\n\n        similar = find_similar_sequences_detailed(seq, train_seqs_df, train_coords_dict, top_n=5)\n        preds   = []\n        used    = set()\n\n        for i, (tmpl_id, tmpl_seq, sim, tmpl_coords, pct_id) in enumerate(similar):\n            if len(preds) >= N_SAMPLE:\n                break\n            if sim < MIN_SIMILARITY or pct_id < MIN_PERCENT_IDENTITY:\n                break           # list is sorted by sim, so no point continuing\n            if tmpl_id in used:\n                continue\n\n            rng     = np.random.default_rng((abs(hash(tid)) + i * 10007) % (2**32))\n            adapted = adapt_template_to_query(seq, tmpl_seq, tmpl_coords)\n\n            # Diversity transforms (same strategy as the 0-409 TBM notebook)\n            slot = len(preds)\n            if slot == 0:\n                X = adapted\n            elif slot == 1:\n                X = adapted + rng.normal(0, max(0.01, (0.40 - sim) * 0.06), adapted.shape)\n            elif slot == 2:\n                longest = max(segs, key=lambda se: se[1] - se[0])\n                X = apply_hinge(adapted, longest, rng)\n            elif slot == 3:\n                X = jitter_chains(adapted, segs, rng)\n            else:\n                X = smooth_wiggle(adapted, segs, rng)\n\n            refined = adaptive_rna_constraints(X, tid, segments_map, confidence=sim)\n            preds.append(refined)\n            used.add(tmpl_id)\n\n        template_predictions[tid] = preds\n        n_needed = N_SAMPLE - len(preds)\n        if n_needed > 0:\n            protenix_queue[tid] = (n_needed, seq)\n            print(f\"  {tid} ({len(seq)} nt): {len(preds)} TBM → need {n_needed} from Protenix\")\n        else:\n            print(f\"  {tid} ({len(seq)} nt): all {N_SAMPLE} from TBM ✓\")\n\n    elapsed = time.time() - t0\n    n_full  = len(test_df) - len(protenix_queue)\n    print(f\"\\nPhase 1 done in {elapsed:.1f}s\")\n    print(f\"  Fully covered by TBM : {n_full}\")\n    print(f\"  Need Protenix        : {len(protenix_queue)}\")\n    return template_predictions, protenix_queue\n\n\n# ─────────────── Protenix Helper Function ───────────────────────────────────────────────────\nimport numpy as np\n\n\n\n\ndef merge_chunks_weighted(chunk_coords_list, chunk_ranges, seq_len, method=\"cosine\", sigma=2.0):\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, chunk_coords_list)):\n        cl = coords.shape[0]\n        ae = min(s + cl, seq_len)\n        ul = ae - s\n        w = np.ones(ul, dtype=np.float64)\n\n        # Overlap with previous\n        if i > 0:\n            ov_s2 = max(s, chunk_ranges[i-1][0])\n            ov_e2 = min(e, chunk_ranges[i-1][1])\n            rl = ov_e2 - ov_s2\n            if rl > 0:\n                if method == \"linear\":\n                    w[:rl] = np.linspace(0., 1., rl)\n                elif method == \"cosine\":\n                    w[:rl] = (1 - np.cos(np.linspace(0, np.pi, rl))) / 2\n                elif method == \"gaussian\":\n                    x = np.arange(rl)\n                    center = rl / 2\n                    w[:rl] = np.exp(-((x - center) ** 2) / (2 * sigma**2))\n\n        # Overlap with next\n        if i < len(chunk_ranges) - 1:\n            ns2 = chunk_ranges[i+1][0]\n            rl = ae - ns2\n            if rl > 0:\n                start = ns2 - s\n                if method == \"linear\":\n                    w[start:start+rl] = np.linspace(1., 0., rl)\n                elif method == \"cosine\":\n                    w[start:start+rl] = (1 + np.cos(np.linspace(0, np.pi, rl))) / 2\n                elif method == \"gaussian\":\n                    x = np.arange(rl)\n                    center = rl / 2\n                    w[start:start+rl] = np.exp(-((x - center) ** 2) / (2 * sigma**2))\n\n        full[s:ae] += coords[:ul] * w[:, None]\n        weights[s:ae] += w\n\n    mask = weights > 0\n    full[mask] /= weights[mask, None]\n    return full\n   \ndef kabsch_align(P, Q):\n    cP, cQ = P.mean(0), Q.mean(0)\n    Pc, Qc = P - cP, Q - cQ\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: S[2,2] = -1\n    R = Vt.T @ S @ U.T\n    return R, cQ - R @ cP\n\ndef stitch_chunks_sequential(chunk_coords_list, chunk_ranges, seq_len):\n    aligned = [chunk_coords_list[0].copy()]\n    for i in range(1, len(chunk_coords_list)):\n        ps, pe = chunk_ranges[i-1]\n        cs, ce = chunk_ranges[i]\n        ov_s, ov_e = cs, min(pe, ce)\n        if ov_e - ov_s < 3:\n            aligned.append(chunk_coords_list[i].copy())\n            continue\n        prev_ov = aligned[i-1][ov_s-ps:ov_e-ps]\n        cur_ov = chunk_coords_list[i][ov_s-cs:ov_e-cs]\n        valid = ~(np.isnan(prev_ov).any(1) | np.isnan(cur_ov).any(1))\n        if valid.sum() < 3:\n            aligned.append(chunk_coords_list[i].copy())\n            continue\n        R, t = kabsch_align(cur_ov[valid], prev_ov[valid])\n        aligned.append((chunk_coords_list[i] @ R.T) + t)\n\n    # merge overlaps with linear weights\n    return merge_chunks_weighted(aligned, chunk_ranges, seq_len, method=\"linear\")\n\ndef stitch_chunks_global(chunk_coords_list, chunk_ranges, seq_len):\n    aligned = [chunk_coords_list[0].copy()]\n    for i in range(1, len(chunk_coords_list)):\n        cs, ce = chunk_ranges[i]\n        fs, fe = chunk_ranges[0]\n        ov_s = max(cs, fs)\n        ov_e = min(ce, fe)\n        if ov_e - ov_s < 3:\n            aligned.append(chunk_coords_list[i].copy())\n            continue\n        prev_ov = aligned[0][ov_s-fs:ov_e-fs]\n        cur_ov = chunk_coords_list[i][ov_s-cs:ov_e-cs]\n        valid = ~(np.isnan(prev_ov).any(1) | np.isnan(cur_ov).any(1))\n        if valid.sum() < 3:\n            aligned.append(chunk_coords_list[i].copy())\n            continue\n        R, t = kabsch_align(cur_ov[valid], prev_ov[valid])\n        aligned.append((chunk_coords_list[i] @ R.T) + t)\n\n    # Merge using smooth cosine weights\n    return merge_chunks_weighted(aligned, chunk_ranges, seq_len, method=\"cosine\")\n\ndef stitch_chunks_graph(chunk_coords_list, chunk_ranges, seq_len, max_iter=5):\n    \"\"\"\n    Deterministic multi-chunk alignment using overlaps.\n    Each chunk aligned to neighbors in sorted order.\n    \"\"\"\n    N = len(chunk_coords_list)\n    aligned = [c.copy() for c in chunk_coords_list]\n\n    # Build neighbor overlaps (sorted for determinism)\n    neighbors = []\n    for i in range(N):\n        ns = []\n        s1, e1 = chunk_ranges[i]\n        for j in range(N):\n            if i == j: continue\n            s2, e2 = chunk_ranges[j]\n            ov_s = max(s1, s2)\n            ov_e = min(e1, e2)\n            if ov_e - ov_s >= 3:\n                ns.append((j, ov_s, ov_e))\n        ns.sort(key=lambda x: x[0])\n        neighbors.append(ns)\n\n    for it in range(max_iter):\n        for i in range(N):\n            for j, ov_s, ov_e in neighbors[i]:\n                s1, e1 = chunk_ranges[i]\n                s2, e2 = chunk_ranges[j]\n                ov_i = aligned[i][ov_s - s1:ov_e - s1]\n                ov_j = aligned[j][ov_s - s2:ov_e - s2]\n                valid = ~(np.isnan(ov_i).any(1) | np.isnan(ov_j).any(1))\n                if valid.sum() < 3: continue\n                R, t = kabsch_align(ov_i[valid], ov_j[valid])\n                aligned[i] = (aligned[i] @ R.T) + t\n\n    return merge_chunks_weighted(aligned, chunk_ranges, seq_len, method=\"cosine\")\n\ndef _get_c1_mask(data, atom_array, chunk_seq_len):\n        \"\"\"\n        v18: Reference-quality C1 mask selection.\n        Priority: atom_array attributes -> feature dict -> tokatom_idx heuristic.\n        \"\"\"\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        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        m11 = (f[\"atom_to_tokatom_idx\"] == 11).bool()\n        m12 = (f[\"atom_to_tokatom_idx\"] == 12).bool()\n        c11, c12 = m11.sum().item(), m12.sum().item()\n        return m11 if abs(c11 - chunk_seq_len) < abs(c12 - chunk_seq_len) else m12\n\ndef _extract_c1_coords(pred, data, atom_array, chunk_seq_len, raw_coords):\n    \"\"\"Select C1' atom coordinates from Protenix output using reference logic.\"\"\"\n    mask = _get_c1_mask(data, atom_array, chunk_seq_len).to(raw_coords.device)\n    coords = raw_coords[:, mask, :].detach().cpu().numpy()\n    if coords.shape[1] > 1:\n        diffs = np.linalg.norm(coords[0,1:]-coords[0,:-1], axis=-1)\n        if np.all(diffs < 1e-4):\n            print(f\"    WARNING: collapsed coordinates detected\"); return None\n    if coords.shape[1] != chunk_seq_len:\n        if coords.shape[1] == 1 and chunk_seq_len > 1: 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    return coords\n\ndef chunk_test_df(test_df, max_len=512, overlap=64):\n    rows = []\n\n    for target_id, (n_needed, full_seq) in test_df.items(): \n      \n        chunks = split_sequence(full_seq, max_len, overlap)\n\n        for i, (start, end, subseq) in enumerate(chunks):\n            rows.append({\n                \"target_id\": f\"{target_id}__chunk{i:03d}\",\n                \"parent_id\": target_id,\n                \"chunk_index\": i,\n                \"start\": start,\n                \"end\": end,\n                \"sequence\": subseq\n            })\n\n    return pd.DataFrame(rows)\n\ndef split_sequence(seq: str, max_len: int = 1024, overlap: int = 128):\n    \"\"\"\n    Split a long sequence into overlapping chunks.\n    Returns list of (start_idx, end_idx, chunk_seq)\n    \"\"\"\n    chunks = []\n    start = 0\n    while start < len(seq):\n        end = min(start + max_len, len(seq))\n        chunks.append((start, end, seq[start:end]))\n        if end == len(seq):\n            break\n        start += max_len - overlap\n    return chunks\n\n# -------------------- Protenix runner (deterministic) --------------------\ndef run_protenix_deterministic(protenix_queue, work_dir, seed_base=SEED):\n    from runner.inference import InferenceRunner, update_inference_configs\n    from protenix.data.inference.infer_dataloader import InferenceDataset\n\n    # Chunk sequences\n    test_df_chunk = chunk_test_df(protenix_queue, max_len=MAX_SEQ_LEN, overlap=CHUNK_OVERLAP)\n    test_df_chunk = test_df_chunk.sort_values(\"start\").reset_index(drop=True)\n    input_json_path = str(work_dir / \"protenix_queue_input.json\")\n    build_input_json(test_df_chunk, input_json_path)\n\n    configs = build_configs(input_json_path, str(work_dir / \"outputs\"), MODEL_NAME)\n    runner  = InferenceRunner(configs)\n    dataset = InferenceDataset(configs)\n\n    raw_predictions = {}\n\n    for i in tqdm(range(len(dataset)), desc=\"Protenix\"):\n        chunk_seed = seed_base + i\n        random.seed(chunk_seed)\n        np.random.seed(chunk_seed)\n        torch.manual_seed(chunk_seed)\n\n        data, atom_array, error_message = dataset[i]\n        chunk_id = data.get(\"sample_name\", f\"sample_{i}\")\n        parent_id = chunk_id.split(\"__chunk\")[0]\n        n_needed, _ = protenix_queue[parent_id]\n        sub_seq_len = data[\"N_token\"].item()\n\n        if error_message:\n            raw_predictions[chunk_id] = None\n            del data, atom_array; gc.collect(); torch.cuda.empty_cache()\n            continue\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            prediction = runner.predict(data)\n            raw_coords = prediction[\"coordinate\"]\n            coords     = _extract_c1_coords(prediction, data, atom_array, sub_seq_len, raw_coords)\n            raw_predictions[chunk_id] = coords\n            del prediction, raw_coords\n        except Exception:\n            raw_predictions[chunk_id] = None\n        finally:\n            del data, atom_array; gc.collect(); torch.cuda.empty_cache()\n\n    # Stitch chunks deterministically\n    protenix_preds = {}\n    for target_id, (n_needed, full_seq) in protenix_queue.items():\n        seq_len = len(full_seq)\n        chunks = test_df_chunk[test_df_chunk[\"parent_id\"] == target_id]\n        if chunks.empty: continue\n        if len(chunks) == 1:\n            cid = chunks.iloc[0][\"target_id\"]\n            protenix_preds[target_id] = raw_predictions.get(cid)\n            continue\n\n        per_sample = {s:[] for s in range(n_needed)}\n        all_ok = True\n        for _, cinfo in chunks.iterrows():\n            ccoords = raw_predictions.get(cinfo[\"target_id\"])\n            if ccoords is None:\n                all_ok = False; break\n            for s_idx in range(n_needed):\n                si = s_idx if s_idx < ccoords.shape[0] else -1\n                per_sample[s_idx].append((ccoords[si], (cinfo[\"start\"], cinfo[\"end\"])))\n        if not all_ok:\n            protenix_preds[target_id] = None\n            continue\n\n        stitched = []\n        for s_idx in range(n_needed):\n            items = per_sample[s_idx]\n            fc = stitch_chunks_graph([c for c,_ in items],[r for _,r in items], seq_len)\n            stitched.append(fc)\n        protenix_preds[target_id] = np.stack(stitched, axis=0)\n\n    return protenix_preds\n\n\nimport hashlib\n\ndef stable_hash(x: str) -> int:\n    return int(hashlib.md5(x.encode()).hexdigest(), 16) % (2**32)\n\n\nfrom sklearn.cluster import AgglomerativeClustering\n\ndef combine_predictions(\n    test_df: pd.DataFrame,\n    template_preds: dict,\n    protenix_preds: dict,\n    segments_map: dict,\n    n_sample: int = FINAL_N_SAMPLE,\n    use_msa: bool = USE_RNA_MSA,\n    seed_base: int = SEED\n):\n    \"\"\"\n    Combines TBM + Protenix + fallback into final submission.\n    Strategy:\n    - Alignment-based scoring\n    - Clustering for diversity\n    - MSA refinement on top candidates\n    \"\"\"\n    all_rows = []\n\n    for _, row in test_df.iterrows():\n        tid = row[\"target_id\"]\n        seq = row[\"sequence\"]\n        # ── Collect candidates ─────────\n        candidates = []\n\n        # TBM predictions\n        tbm_preds = template_preds.get(tid, [])\n        candidates.extend(tbm_preds)\n\n        # Protenix predictions\n        ptx_raw = protenix_preds.get(tid)\n        if ptx_raw is not None and ptx_raw.ndim == 3:\n            candidates.extend([ptx_raw[i] for i in range(ptx_raw.shape[0])])\n            \n        N = len(candidates)\n        \n        # ── Fallback if nothing exists ─\n        if N < n_sample:\n            print(f\"[WARNING] No predictions for {tid}, using de-novo fallback\")\n            for i in range(n_sample-N):\n                seed_val = stable_hash(tid) + i * 10007 + seed_base\n                dn = generate_rna_structure(seq, seed=seed_val)\n                dn = adaptive_rna_constraints(dn, tid, segments_map, confidence=0.05)\n                candidates.append(dn)\n                \n        # ── Alignment-based scoring ───\n        dist_matrix = np.zeros((N, N))\n        for i in range(N):\n            for j in range(i + 1, N):\n                try:\n                    R, t = kabsch_align(candidates[i], candidates[j])\n                    aligned = (candidates[i] @ R.T) + t\n                    d = np.mean(np.linalg.norm(aligned - candidates[j], axis=1))\n                    dist_matrix[i, j] = dist_matrix[j, i] = d\n                except:\n                    dist_matrix[i, j] = dist_matrix[j, i] = 1e6  # fallback for errors\n\n        # ── Clustering for diversity ───\n        if N > n_sample:\n            clustering = AgglomerativeClustering(\n                n_clusters=n_sample, metric='precomputed', linkage='average'\n            )\n            cluster_labels = clustering.fit_predict(dist_matrix)\n            selected = []\n            for c in range(n_sample):\n                idxs = np.where(cluster_labels == c)[0]\n                if len(idxs) == 1:\n                    selected.append(idxs[0])\n                else:\n                    mean_dists = dist_matrix[np.ix_(idxs, idxs)].mean(1)\n                    selected.append(idxs[np.argmin(mean_dists)])\n            candidates = [candidates[i] for i in selected]\n        else:\n            candidates = candidates[:n_sample]\n\n        # ── Optional: MSA / segment refinement ───\n        if use_msa:\n            for i in range(len(candidates)):\n                try:\n                    candidates[i] = adaptive_rna_constraints(\n                        candidates[i],\n                        tid,\n                        segments_map,\n                        confidence=0.8,\n                        passes=2\n                    )\n                except:\n                    pass\n\n        # ── Stack and convert to submission rows ──\n        stacked = np.stack(candidates[:n_sample], axis=0)\n        all_rows.extend(coords_to_rows(tid, seq, stacked))\n\n    # ── Build submission DataFrame ─────────\n    cols = [\"ID\", \"resname\", \"resid\"] + [\n        f\"{c}_{i}\" for i in range(1, n_sample + 1)\n        for c in [\"x\", \"y\", \"z\"]\n    ]\n    sub = pd.DataFrame(all_rows)\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\n    return sub[cols]\n\n\n\ndef main() -> None:\n    test_csv, output_csv, code_dir, root_dir = resolve_paths()\n\n    if not os.path.isdir(code_dir):\n        raise FileNotFoundError(\n            f\"Missing PROTENIX_CODE_DIR: {code_dir}. \"\n            \"Set PROTENIX_CODE_DIR to the repo path.\"\n        )\n\n    os.environ[\"PROTENIX_ROOT_DIR\"] = root_dir\n    sys.path.append(code_dir)\n    ensure_required_files(root_dir)\n    seed_everything(SEED)\n\n    # ── Load test data ──────────────────────────────────────────────────────\n    test_df_full = pd.read_csv(test_csv)\n    if DO_SINGLE_TEST:\n        test_df = (test_df_full[test_df_full.target_id == SINGLE_TEST] if not IS_KAGGLE\n                else test_df_full).reset_index(drop=True)\n    else:\n        test_df      = (test_df_full.head(LOCAL_N_SAMPLES) if not IS_KAGGLE\n                        else test_df_full).reset_index(drop=True)\n        \n    print(f\"Test targets : {len(test_df)}\"\n          + (\" (LOCAL MODE)\" if not IS_KAGGLE else \"\"))\n\n    # ── Load training data for TBM ──────────────────────────────────────────\n    print(\"\\nLoading training data for TBM …\")\n    train_seqs   = pd.read_csv(DEFAULT_TRAIN_CSV)\n    val_seqs     = pd.read_csv(DEFAULT_VAL_CSV)\n    train_labels = pd.read_csv(DEFAULT_TRAIN_LBLS)\n    val_labels   = pd.read_csv(DEFAULT_VAL_LBLS)\n\n    combined_seqs   = pd.concat([train_seqs, val_seqs], ignore_index=True)\n    combined_labels = pd.concat([train_labels, val_labels], ignore_index=True)\n    train_coords    = process_labels(combined_labels)\n    segments_map, _ = build_segments_map(test_df)\n\n    print(f\"Template pool: {len(combined_seqs)} sequences, {len(train_coords)} structures\")\n\n    # ─── PHASE 1: TBM ──────────────────────────────────────────────────────\n    template_preds, protenix_queue = tbm_phase(\n        test_df, combined_seqs, train_coords, segments_map\n    )\n\n    # free RAM\n    del train_seqs, val_seqs, train_labels, val_labels, combined_seqs, combined_labels, train_coords, test_df_full\n    import gc\n    gc.collect()\n\n    # ─── PHASE 2: Protenix ───────────────────────────────────────────────\n   \n    if protenix_queue and USE_PROTENIX:\n        print(f\"\\n{'='*60}\")\n        print(f\"PHASE 2: Protenix for {len(protenix_queue)} targets\")\n        print(f\"{'='*60}\")\n\n        work_dir = Path(\"/kaggle/working\")\n        work_dir.mkdir(parents=True, exist_ok=True)\n        protenix_preds = run_protenix_deterministic(protenix_queue, work_dir, seed_base=101)\n    elif protenix_queue and not USE_PROTENIX:\n        print(f\"\\nPHASE 2 skipped (USE_PROTENIX=False). \"\n              f\"De-novo fallback will cover {len(protenix_queue)} targets.\")\n\n    # ─── PHASE 3: Combine TBM + Protenix + fallback ───────────────────────\n    print(f\"\\n{'='*60}\")\n    print(\"PHASE 3: Combine TBM + Protenix + de-novo fallback\")\n    print(f\"{'='*60}\")\n\n\n    combine_output=combine_predictions(test_df,template_preds,protenix_preds,segments_map)\n    print(combine_output.head())\n    print(combine_output.shape)\n    print(combine_output.isna().sum().sum(), \"NaNs\")\n    \n    # Check columns\n    print(combine_output.columns.tolist())\n    \n    # Check rows\n    expected_rows = test_df[\"sequence\"].str.len().sum()\n    print(\"Expected:\", expected_rows, \"Actual:\", len(combine_output))\n    # ── Save submission ───────────────────────────────────────────────────\n    combine_output.to_csv(output_csv, index=False)\n    print(f\"\\n✓ Saved submission to {output_csv}  ({len(combine_output):,} rows)\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:31:20.590485Z","iopub.execute_input":"2026-03-25T09:31:20.590854Z","iopub.status.idle":"2026-03-25T09:31:21.961861Z","shell.execute_reply.started":"2026-03-25T09:31:20.590832Z","shell.execute_reply":"2026-03-25T09:31:21.961298Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"\nif __name__ == \"__main__\":\n    main()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:31:21.962767Z","iopub.execute_input":"2026-03-25T09:31:21.963206Z","iopub.status.idle":"2026-03-25T09:35:03.2825Z","shell.execute_reply.started":"2026-03-25T09:31:21.96317Z","shell.execute_reply":"2026-03-25T09:35:03.281631Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#read submission.csv\nsubmission_path = \"/kaggle/working/submission.csv\"\nsubmission_df = pd.read_csv(submission_path)\nprint(submission_df.head(5))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:35:03.284401Z","iopub.execute_input":"2026-03-25T09:35:03.284715Z","iopub.status.idle":"2026-03-25T09:35:03.297449Z","shell.execute_reply.started":"2026-03-25T09:35:03.28469Z","shell.execute_reply":"2026-03-25T09:35:03.296713Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if not IS_KAGGLE:\n    RESULT_ID = 0\n    import pandas as pd\n    import numpy as np\n    import matplotlib.pyplot as plt\n    from mpl_toolkits.mplot3d import Axes3D\n\n   \n\n   # ---------------- Helper: extract coordinates from a row ----------------\n    def row_to_coords(row):\n        coords_list = []\n        for i in range(1, 6):  # x_1 ... x_5\n            coords_list.append([row[f\"x_{i}\"], row[f\"y_{i}\"], row[f\"z_{i}\"]])\n        return np.array(coords_list, dtype=np.float32)  # shape: [5, 3]\n    \n    # --------------- Build predictions list for a target --------------------\n    def extract_preds_from_csv(df, target_id, n_sample=5):\n        \"\"\"\n        Returns a list of n_sample arrays, each of shape [seq_len, 3].\n        Each \"prediction\" corresponds to one of the 5 coordinates per residue in the CSV.\n        \"\"\"\n        target_rows = df[df[\"ID\"].str.split(\"_\").str[0] == target_id].sort_values(\"resid\").reset_index(drop=True)\n        \n        seq_len = len(target_rows)\n        preds = [np.zeros((seq_len, 3), dtype=np.float32) for _ in range(n_sample)]\n        \n        for residue_idx, row in target_rows.iterrows():  # <-- use iterrows\n            coords = row_to_coords(row)  # [5, 3]\n            for i in range(n_sample):\n                preds[i][residue_idx] = coords[i]  # take i-th coord for prediction i\n        \n        return preds\n\n    # ---------------- Example: pick a sample target ------------------------\n    df = pd.read_csv(\"/kaggle/working/submission.csv\")\n    if DO_SINGLE_TEST:\n        print(\"iii\")\n        sample_preds = extract_preds_from_csv(df, SINGLE_TEST)\n    else :\n        print(\"fdsf\")\n        test_df = pd.read_csv(DEFAULT_TEST_CSV)\n        sample_target_id = test_df[\"target_id\"][RESULT_ID]  # third unique target\n        sample_preds = extract_preds_from_csv(df, sample_target_id)\n\n    # ---------------- 3D plot of predictions -------------------------------\n    fig = plt.figure(figsize=(16, 4))\n    for i in range(len(sample_preds)):\n        ax = fig.add_subplot(1, len(sample_preds), i+1, projection='3d')\n        coords = sample_preds[i]\n        ax.plot(coords[:, 0], coords[:, 1], coords[:, 2], 'o-', markersize=3, linewidth=1)\n        ax.set_title(f'Prediction {i+1}')\n        ax.set_xlabel('X')\n        ax.set_ylabel('Y')\n        ax.set_zlabel('Z')\n    plt.suptitle(f'Ensemble Predictions: {SINGLE_TEST} (Length: {sample_preds[0].shape[0]})',\n                 fontweight='bold')\n    plt.tight_layout()\n    plt.show()\n\n    # ---------------- Distance analysis -----------------------------------\n    all_distances = []\n    fig, axes = plt.subplots(1, 2, figsize=(14, 4))\n\n    # Sequential distances\n    for pred in sample_preds:\n        dists = [np.linalg.norm(pred[i+1] - pred[i]) for i in range(len(pred)-1)]\n        axes[0].plot(range(len(dists)), dists, alpha=0.6, linewidth=1)\n        all_distances.extend(dists)\n\n    axes[0].axhline(y=6.0, color='red', linestyle='--', linewidth=1.5, label='Target: 6.0Å')\n    axes[0].set_xlabel('Residue Index')\n    axes[0].set_ylabel('Distance (Å)')\n    axes[0].set_title('Sequential Residue Distances')\n    axes[0].legend()\n    axes[0].grid(True, alpha=0.3)\n\n    # Histogram of distances\n    axes[1].hist(all_distances, bins=30, color='steelblue', edgecolor='black', alpha=0.7)\n    axes[1].axvline(x=6.0, color='red', linestyle='--', linewidth=2, label='Target: 6.0Å', alpha=0.3)\n    axes[1].set_xlabel('Distance (Å)')\n    axes[1].set_ylabel('Frequency')\n    axes[1].set_title('Distance Distribution')\n    axes[1].legend()\n    axes[1].grid(True, alpha=0.3, axis='y')\n\n    plt.tight_layout()\n    plt.show()\n\n    print(f\"Mean distance: {np.mean(all_distances):.2f}Å | Std: {np.std(all_distances):.2f}Å\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-03-25T09:35:03.298504Z","iopub.execute_input":"2026-03-25T09:35:03.298753Z","iopub.status.idle":"2026-03-25T09:35:04.220794Z","shell.execute_reply.started":"2026-03-25T09:35:03.298731Z","shell.execute_reply":"2026-03-25T09:35:04.220211Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}