{"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}],"dockerImageVersionId":30919,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# DO UPVOTE GUYS\n","metadata":{}},{"cell_type":"markdown","source":"## BASLINE CREDIT GOES TO ITS RESPECTIVE AUTHOR\n# Template + Protenix Combined RNA 3D Folding\n\n## Strategy\n\n1. **Template-Based Modeling** (primary): Find similar sequences in training data and adapt their known 3D structures\n\n    Ref: https://www.kaggle.com/code/jakupymeraj/rna-folding-prediction-12feb-v1\n\n2. **Protenix** (fallback): When good templates aren't available (similarity < 0.0 or percent identity < 50%), use Protenix\n\n    Ref: https://www.kaggle.com/code/zoushuxian/aido-rna-augmented-protenix\n\n### Pipeline\n\n```\nPhase 1: Template matching for ALL targets \n  -> Good templates found?  -> Use adapted coordinates\n  -> Poor/no templates?     -> Queue for Protenix\n\nPhase 2: Protenix inference for queued targets \n  -> Generate N predictions per target (only as many as needed)\n\nPhase 3: Combine template + Protenix predictions\n  -> 5 predictions per target for submission\n```","metadata":{}},{"cell_type":"code","source":"import os\nimport sys\n\n# Set data path (Kaggle or local)\nif os.path.exists('/kaggle/input/stanford-rna-3d-folding-2/'):\n    IS_KAGGLE = True\n    DATA_PATH = '/kaggle/input/stanford-rna-3d-folding-2/'\n    OUTPUT_PATH = '/kaggle/working/output'\n    USALIGN_BIN = '/kaggle/working/USalign'\n    PROTENIX_DIR = '/kaggle/working/Protenix'\n    ! cp /kaggle/input/protenix-packages/packages/USalign /kaggle/working/\n    ! chmod +x /kaggle/working/USalign\n    sys.path.insert(0, '/kaggle/input/rna-3d-utils/')\nelif os.path.exists('/data/kaggle/stanford-rna-3d-folding-2/'):\n    IS_KAGGLE = False\n    DATA_PATH = '/data/kaggle/stanford-rna-3d-folding-2/'\n    OUTPUT_PATH = './output'\n    USALIGN_BIN = './USalign/USalign'\n    PROTENIX_DIR = f'{DATA_PATH}/protenix_kaggle'\n    sys.path.insert(0, os.path.join(os.getcwd(), 'kaggle_util_ds'))\nelse:\n    raise ValueError(\"Data path not found\")\n\nprint(f\"Data path: {DATA_PATH}\")\nprint(f\"Protenix dir: {PROTENIX_DIR}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if IS_KAGGLE:\n    !cp -r /kaggle/input/protenix-packages/packages /kaggle/working\n    %cd /kaggle/working/packages\n    !pip install --no-deps --exists-action=i *.whl\n    %cd /kaggle/working\n\n    !mv /kaggle/working/packages/ihm-2.3/ihm-2.3 /kaggle/working\n    !mv /kaggle/working/packages/modelcif-0.7/modelcif-0.7 /kaggle/working\n\n    !pip install /kaggle/working/ihm-2.3\n    !pip install /kaggle/working/modelcif-0.7\n\n    !rm -rf /kaggle/working/ihm-2.3\n    !rm -rf /kaggle/working/modelcif-0.7\n\n    !pip install /kaggle/input/biopython/biopython-1.85-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl\n    !pip install /kaggle/input/ml-collections/ml_collections-1.0.0-py3-none-any.whl\n\n    !rm -rf /kaggle/working/packages\n\n    !cp -r /kaggle/input/protenix-mg-packages/protenix_mg_packages /kaggle/working\n    %cd /kaggle/working/protenix_mg_packages\n    !pip install --no-deps --exists-action=i *.whl\n    %cd /kaggle/working\n    !rm -rf /kaggle/working/protenix_mg_packages\n\n    !cp -R /kaggle/input/protenix-rmsa-repo/protenix_kaggle /kaggle/working/\n    !mv protenix_kaggle Protenix","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import json\nimport numpy as np\nimport pandas as pd\nimport time\nimport random\nimport warnings\nimport contextlib\nfrom pathlib import Path\n\nwarnings.filterwarnings('ignore')\n\ndef seed_everything(seed: int = 42):\n    random.seed(seed)\n    np.random.seed(seed)\n\nseed_everything(42)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Loading sequence data...\")\ntrain_seqs = pd.read_csv(DATA_PATH + 'train_sequences.csv')\nvalidation_seqs = pd.read_csv(DATA_PATH + 'validation_sequences.csv')\ntest_seqs = pd.read_csv(DATA_PATH + 'test_sequences.csv')\n\ntrain_labels = pd.read_csv(DATA_PATH + 'train_labels.csv')\nvalidation_labels = pd.read_csv(DATA_PATH + 'validation_labels.csv')\n\nprint(f\"Loaded {len(train_seqs)} training sequences\")\nprint(f\"Loaded {len(validation_seqs)} validation sequences\")\nprint(f\"Loaded {len(test_seqs)} test sequences\")\n\nSHOW_VALIDATION = False\nSHOW_ALIGNMENT_DETAILS = True\nMAKE_SUBMISSION = True\nUSE_PROTENIX = True  # Set False to fall back to simple de novo\n\n# Quality thresholds for template selection\nMIN_SIMILARITY = 0.0         # use -1 to disable threshold\nMIN_PERCENT_IDENTITY = 50   # use -1 to disable threshold\n\nDEBUG = False\nif DEBUG:\n    MIN_SIMILARITY = 0.3","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Functions for template modelling","metadata":{}},{"cell_type":"code","source":"\"\"\"\nTemplate-based modelling for RNA 3D structure prediction.\n\nRef:https://www.kaggle.com/code/jakupymeraj/rna-folding-prediction-12feb-v1\n\nImplements an optimized template-based modeling (TBM) approach:\n1. Expanded Template Pool: Combines training + validation data for better template coverage\n2. Optimized Alignment: Uses Bio.Align.PairwiseAligner with tuned gap penalties\n3. Adaptive Constraints: Segment-aware physics-based constraints scaled by confidence\n4. Ensemble Diversity: Generates 5 predictions with hinge, jitter, and wiggle transforms\n\"\"\"\n\nimport numpy as np\nimport pandas as pd\nfrom Bio.Align import PairwiseAligner\n\n\n# ---------------------------------------------------------------------------\n# Aligner setup\n# ---------------------------------------------------------------------------\n\ndef make_aligner():\n    \"\"\"Create a PairwiseAligner with tuned gap penalties for RNA TBM.\"\"\"\n    al = PairwiseAligner()\n    al.mode = 'global'\n    al.match_score = 2\n    al.mismatch_score = -1.5\n\n    al.open_gap_score = -8\n    al.extend_gap_score = -0.4\n\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\n# ---------------------------------------------------------------------------\n# FASTA / stoichiometry / segment helpers\n# ---------------------------------------------------------------------------\n\ndef parse_fasta(fasta_content: str):\n    \"\"\"Parse FASTA string to {chain_id: sequence_string}.\"\"\"\n    out = {}\n    cur = None\n    seq_parts = []\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(seq_parts)\n            cur = line[1:].split()[0]\n            seq_parts = []\n        else:\n            seq_parts.append(line.replace(\" \", \"\"))\n    if cur is not None:\n        out[cur] = \"\".join(seq_parts)\n    return out\n\n\ndef parse_stoichiometry(stoich: str):\n    if pd.isna(stoich) or str(stoich).strip() == \"\":\n        return []\n    out = []\n    for part in str(stoich).split(';'):\n        ch, cnt = part.split(':')\n        out.append((ch.strip(), int(cnt)))\n    return out\n\n\ndef get_chain_segments(row):\n    \"\"\"\n    Returns list of (start, end) segments in row['sequence'] corresponding\n    to chain copies in stoichiometry order.\n    Falls back to single segment if parsing fails.\n    \"\"\"\n    seq = row['sequence']\n    stoich = row.get('stoichiometry', '')\n    all_seq = row.get('all_sequences', '')\n\n    if pd.isna(stoich) or pd.isna(all_seq) or str(stoich).strip() == \"\" or str(all_seq).strip() == \"\":\n        return [(0, len(seq))]\n\n    try:\n        chain_dict = parse_fasta(all_seq)\n        order = parse_stoichiometry(stoich)\n        segs = []\n        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                L = len(base)\n                segs.append((pos, pos + L))\n                pos += L\n        if pos != len(seq):\n            return [(0, len(seq))]\n        return segs\n    except Exception:\n        return [(0, len(seq))]\n\n\ndef build_segments_map(df):\n    seg_map = {}\n    stoich_map = {}\n    for _, r in df.iterrows():\n        tid = r['target_id']\n        seg_map[tid] = get_chain_segments(r)\n        stoich_map[tid] = str(r.get('stoichiometry', '') if not pd.isna(r.get('stoichiometry', '')) else '')\n    return seg_map, stoich_map\n\n\n# ---------------------------------------------------------------------------\n# Data processing\n# ---------------------------------------------------------------------------\n\ndef process_labels(labels_df):\n    \"\"\"\n    Convert labels dataframe to dictionary mapping target_id -> 3D coordinates.\n\n    Args:\n        labels_df: DataFrame with columns [ID, resid, x_1, y_1, z_1, ...]\n\n    Returns:\n        dict: {target_id: np.array([[x1,y1,z1], [x2,y2,z2], ...])}\n    \"\"\"\n    coords_dict = {}\n    prefixes = labels_df['ID'].str.rsplit('_', n=1).str[0]\n    for id_prefix, group in labels_df.groupby(prefixes):\n        coords_dict[id_prefix] = group.sort_values('resid')[['x_1', 'y_1', 'z_1']].values\n    return coords_dict\n\n\n# ---------------------------------------------------------------------------\n# Template search\n# ---------------------------------------------------------------------------\n\ndef find_similar_sequences(query_seq, train_seqs_df, train_coords_dict, top_n=5):\n    \"\"\"\n    Find the most similar template sequences using PairwiseAligner.score()\n    for fast scoring without traceback overhead.\n\n    Returns:\n        List of (target_id, sequence, normalized_score, coordinates) tuples.\n    \"\"\"\n    similar_seqs = []\n\n    for _, row in train_seqs_df.iterrows():\n        target_id, train_seq = row['target_id'], row['sequence']\n        if target_id not in train_coords_dict:\n            continue\n\n        if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > 0.3:\n            continue\n\n        raw_score = _aligner.score(query_seq, train_seq)\n        normalized_score = raw_score / (2 * min(len(query_seq), len(train_seq)))\n        similar_seqs.append((target_id, train_seq, normalized_score, train_coords_dict[target_id]))\n\n    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\n\ndef _build_aligned_strings(query_seq, template_seq, alignment):\n    \"\"\"Build gapped alignment strings from alignment.aligned segments.\"\"\"\n    q_segments, t_segments = alignment.aligned\n    aligned_q = []\n    aligned_t = []\n    qi = 0\n    ti = 0\n\n    for (qs, qe), (ts, te) in zip(q_segments, t_segments):\n        while qi < qs:\n            aligned_q.append(query_seq[qi])\n            aligned_t.append('-')\n            qi += 1\n        while ti < ts:\n            aligned_q.append('-')\n            aligned_t.append(template_seq[ti])\n            ti += 1\n        for q_pos, t_pos in zip(range(qs, qe), range(ts, te)):\n            aligned_q.append(query_seq[q_pos])\n            aligned_t.append(template_seq[t_pos])\n        qi = qe\n        ti = te\n\n    while qi < len(query_seq):\n        aligned_q.append(query_seq[qi])\n        aligned_t.append('-')\n        qi += 1\n    while ti < len(template_seq):\n        aligned_q.append('-')\n        aligned_t.append(template_seq[ti])\n        ti += 1\n\n    return ''.join(aligned_q), ''.join(aligned_t)\n\n\ndef find_similar_sequences_detailed(query_seq, train_seqs_df, train_coords_dict,\n                                    temporal_cutoff=None, top_n=5):\n    \"\"\"\n    Find similar template sequences with full alignment details.\n\n    Unlike find_similar_sequences (which uses score-only for speed), this\n    performs full alignment to compute percent identity and aligned strings.\n\n    Returns:\n        List of (target_id, sequence, normalized_score, coordinates,\n                 percent_identity, aligned_query, aligned_template) tuples.\n    \"\"\"\n    similar_seqs = []\n\n    if temporal_cutoff is not None:\n        filtered = train_seqs_df[train_seqs_df['temporal_cutoff'] < temporal_cutoff]\n    else:\n        filtered = train_seqs_df\n\n    for _, row in filtered.iterrows():\n        target_id, train_seq = row['target_id'], row['sequence']\n        if target_id not in train_coords_dict:\n            continue\n\n        if abs(len(train_seq) - len(query_seq)) / max(len(train_seq), len(query_seq)) > 0.3:\n            continue\n\n        alignment = next(iter(_aligner.align(query_seq, train_seq)))\n        raw_score = alignment.score\n        normalized_score = raw_score / (2 * min(len(query_seq), len(train_seq)))\n\n        identical = 0\n        for (qs, qe), (ts, te) in zip(*alignment.aligned):\n            for q_pos, t_pos in zip(range(qs, qe), range(ts, te)):\n                if query_seq[q_pos] == train_seq[t_pos]:\n                    identical += 1\n        percent_identity = 100 * identical / len(query_seq)\n\n        aligned_query, aligned_template = _build_aligned_strings(\n            query_seq, train_seq, alignment\n        )\n\n        similar_seqs.append((\n            target_id, train_seq, normalized_score,\n            train_coords_dict[target_id], percent_identity,\n            aligned_query, aligned_template\n        ))\n\n    similar_seqs.sort(key=lambda x: x[2], reverse=True)\n    return similar_seqs[:top_n]\n\n\n# ---------------------------------------------------------------------------\n# Coordinate adaptation\n# ---------------------------------------------------------------------------\n\ndef adapt_template_to_query(query_seq, template_seq, template_coords):\n    \"\"\"\n    Adapt template 3D coordinates to match query sequence via alignment.\n    Uses alignment.aligned for vectorized coordinate mapping.\n    \"\"\"\n    alignment = next(iter(_aligner.align(query_seq, template_seq)))\n\n    new_coords = np.full((len(query_seq), 3), np.nan)\n\n    for (q_start, q_end), (t_start, t_end) in zip(*alignment.aligned):\n        t_chunk = template_coords[t_start:t_end]\n        if len(t_chunk) == (q_end - q_start):\n            new_coords[q_start:q_end] = t_chunk\n\n    for i in range(len(new_coords)):\n        if np.isnan(new_coords[i, 0]):\n            prev_v = next((j for j in range(i - 1, -1, -1) if not np.isnan(new_coords[j, 0])), -1)\n            next_v = next((j for j in range(i + 1, len(new_coords)) if not np.isnan(new_coords[j, 0])), -1)\n            if prev_v >= 0 and next_v >= 0:\n                w = (i - prev_v) / (next_v - prev_v)\n                new_coords[i] = (1 - w) * new_coords[prev_v] + w * new_coords[next_v]\n            elif prev_v >= 0:\n                new_coords[i] = new_coords[prev_v] + [3, 0, 0]\n            elif next_v >= 0:\n                new_coords[i] = new_coords[next_v] + [3, 0, 0]\n            else:\n                new_coords[i] = [i * 3, 0, 0]\n\n    return np.nan_to_num(new_coords)\n\n\n# ---------------------------------------------------------------------------\n# Geometric constraints (segment-aware, vectorized)\n# ---------------------------------------------------------------------------\n\ndef adaptive_rna_constraints(coordinates, target_id, segments_map, confidence=1.0, passes=2):\n    \"\"\"\n    Apply RNA geometric constraints within each chain segment.\n    Vectorized bond, angle, Laplacian smoothing, and self-avoidance.\n\n    Args:\n        coordinates: (N, 3) array of residue positions\n        target_id: target identifier for segment lookup\n        segments_map: dict mapping target_id -> list of (start, end) tuples\n        confidence: template similarity score (higher = gentler corrections)\n        passes: number of refinement passes\n    \"\"\"\n    coords = coordinates.copy()\n    segments = segments_map.get(target_id, [(0, len(coords))])\n\n    strength = 0.75 * (1.0 - min(confidence, 0.97))\n    strength = max(strength, 0.02)\n\n    for _ in range(passes):\n        for (s, e) in segments:\n            X = coords[s:e]\n            L = e - s\n            if L < 3:\n                coords[s:e] = X\n                continue\n\n            # (1) bond i,i+1 to ~5.95A\n            d = X[1:] - X[:-1]\n            dist = np.linalg.norm(d, axis=1) + 1e-6\n            target = 5.95\n            scale = (target - dist) / dist\n            adj = (d * scale[:, None]) * (0.22 * strength)\n            X[:-1] -= adj\n            X[1:] += adj\n\n            # (2) soft i,i+2 to ~10.2A\n            d2 = X[2:] - X[:-2]\n            dist2 = np.linalg.norm(d2, axis=1) + 1e-6\n            target2 = 10.2\n            scale2 = (target2 - dist2) / dist2\n            adj2 = (d2 * scale2[:, None]) * (0.10 * strength)\n            X[:-2] -= adj2\n            X[2:] += adj2\n\n            # (3) Laplacian smoothing\n            lap = 0.5 * (X[:-2] + X[2:]) - X[1:-1]\n            X[1:-1] += (0.06 * strength) * lap\n\n            # (4) self-avoidance\n            if L >= 25:\n                k = min(L, 160) if L > 220 else L\n                if k < L:\n                    idx = np.linspace(0, L - 1, k).astype(int)\n                else:\n                    idx = np.arange(L)\n\n                P = X[idx]\n                diff = P[:, None, :] - P[None, :, :]\n                distm = np.linalg.norm(diff, axis=2) + 1e-6\n                sep = np.abs(idx[:, None] - idx[None, :])\n\n                mask = (sep > 2) & (distm < 3.2)\n                if np.any(mask):\n                    force = (3.2 - distm) / distm\n                    vec = (diff * force[:, :, None] * mask[:, :, None]).sum(axis=1)\n                    X[idx] += (0.015 * strength) * vec\n\n            coords[s:e] = X\n\n    return coords\n\n\n# ---------------------------------------------------------------------------\n# Diversity transforms\n# ---------------------------------------------------------------------------\n\ndef _rotmat(axis, ang):\n    axis = np.asarray(axis, float)\n    axis = axis / (np.linalg.norm(axis) + 1e-12)\n    x, y, z = axis\n    c, s = np.cos(ang), np.sin(ang)\n    C = 1.0 - c\n    return np.array([\n        [c + x*x*C,     x*y*C - z*s, x*z*C + y*s],\n        [y*x*C + z*s,   c + y*y*C,   y*z*C - x*s],\n        [z*x*C - y*s,   z*y*C + x*s, c + z*z*C]\n    ], dtype=float)\n\n\ndef apply_hinge(coords, seg, rng, max_angle_deg=25):\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    axis = rng.normal(size=3)\n    ang = np.deg2rad(float(rng.uniform(-max_angle_deg, max_angle_deg)))\n    R = _rotmat(axis, ang)\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, segments, rng, max_angle_deg=12, max_trans=1.5):\n    X = coords.copy()\n    global_center = X.mean(axis=0, keepdims=True)\n    for (s, e) in segments:\n        axis = rng.normal(size=3)\n        ang = np.deg2rad(float(rng.uniform(-max_angle_deg, max_angle_deg)))\n        R = _rotmat(axis, ang)\n        shift = rng.normal(size=3)\n        shift = shift / (np.linalg.norm(shift) + 1e-12) * float(rng.uniform(0.0, max_trans))\n        c = X[s:e].mean(axis=0, keepdims=True)\n        X[s:e] = (X[s:e] - c) @ R.T + c + shift\n    X -= X.mean(axis=0, keepdims=True) - global_center\n    return X\n\n\ndef smooth_wiggle(coords, segments, rng, amp=0.8):\n    X = coords.copy()\n    for (s, e) in segments:\n        L = e - s\n        if L < 20:\n            continue\n        n_ctrl = 6\n        ctrl_x = np.linspace(0, L - 1, n_ctrl)\n        ctrl_disp = rng.normal(0, amp, size=(n_ctrl, 3))\n        t = np.arange(L)\n        disp = np.vstack([np.interp(t, ctrl_x, ctrl_disp[:, k]) for k in range(3)]).T\n        X[s:e] += disp\n    return X\n\n\n# ---------------------------------------------------------------------------\n# Main prediction pipeline\n# ---------------------------------------------------------------------------\n\ndef predict_rna_structures(row, train_seqs_df, train_coords_dict, segments_map, n_predictions=5):\n    \"\"\"\n    Generate n_predictions diverse 3D structure predictions for a test RNA.\n\n    Args:\n        row: DataFrame row with 'target_id' and 'sequence'\n        train_seqs_df: training sequences DataFrame\n        train_coords_dict: {target_id: coords} from process_labels\n        segments_map: {target_id: [(start, end), ...]} from build_segments_map\n        n_predictions: number of predictions (default 5)\n\n    Returns:\n        List of n_predictions numpy arrays, each (seq_len, 3).\n    \"\"\"\n    tid = row['target_id']\n    seq = row['sequence']\n    segments = segments_map.get(tid, [(0, len(seq))])\n\n    cands = find_similar_sequences(\n        query_seq=seq, train_seqs_df=train_seqs_df,\n        train_coords_dict=train_coords_dict, top_n=30\n    )\n\n    predictions = []\n    used = set()\n\n    for i in range(n_predictions):\n        seed = (abs(hash(tid)) + i * 10007) % (2**32)\n        rng = np.random.default_rng(seed)\n\n        if not cands:\n            coords = np.zeros((len(seq), 3), dtype=float)\n            for (s, e) in segments:\n                for j in range(s + 1, e):\n                    coords[j] = coords[j - 1] + [5.95, 0, 0]\n            predictions.append(coords)\n            continue\n\n        if i == 0:\n            t_id, t_seq, sim, t_coords = cands[0]\n        else:\n            K = min(12, len(cands))\n            sims = np.array([cands[k][2] for k in range(K)], float)\n            w = np.exp((sims - sims.max()) / 0.08)\n            for k in range(K):\n                if cands[k][0] in used:\n                    w[k] *= 0.10\n            w = w / (w.sum() + 1e-12)\n            k = int(rng.choice(np.arange(K), p=w))\n            t_id, t_seq, sim, t_coords = cands[k]\n\n        used.add(t_id)\n\n        adapted = adapt_template_to_query(\n            query_seq=seq, template_seq=t_seq, template_coords=t_coords\n        )\n\n        if i == 0:\n            X = adapted\n        elif i == 1:\n            X = adapted + rng.normal(0, max(0.01, (0.40 - sim) * 0.06), adapted.shape)\n        elif i == 2:\n            longest = max(segments, key=lambda se: se[1] - se[0])\n            X = apply_hinge(adapted, longest, rng, max_angle_deg=22)\n        elif i == 3:\n            X = jitter_chains(adapted, segments, rng, max_angle_deg=10, max_trans=1.0)\n        else:\n            X = smooth_wiggle(adapted, segments, rng, amp=0.7)\n\n        refined = adaptive_rna_constraints(X, tid, segments_map, confidence=sim, passes=2)\n        predictions.append(refined)\n\n    return predictions\n\n\n# ---------------------------------------------------------------------------\n# De novo fallback\n# ---------------------------------------------------------------------------\n\ndef generate_rna_structure(sequence, seed=None):\n    \"\"\"Generate idealized A-form RNA helix as de novo fallback.\"\"\"\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        angle = i * 0.6\n        coords[i] = [10.0 * np.cos(angle), 10.0 * np.sin(angle), i * 2.5]\n    return coords","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(\"Processing coordinates...\")\ntrain_coords_dict = process_labels(train_labels)\n\ncombined_seqs = pd.concat([train_seqs, validation_seqs], ignore_index=True)\ncombined_labels = pd.concat([train_labels, validation_labels], ignore_index=True)\ncombined_coords_dict = process_labels(combined_labels)\n\nprint(f\"Processed {len(train_coords_dict)} training structures\")\nprint(f\"Processed {len(combined_coords_dict)} combined (train+val) structures\")\n\n# Build segment maps for validation and test targets\nvalidation_segments_map, _ = build_segments_map(validation_seqs)\ntest_segments_map, _ = build_segments_map(test_seqs)\n\nprint(f\"Built segment maps: {len(validation_segments_map)} validation, {len(test_segments_map)} test\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Protenix Utilities\n\nFunctions for running Protenix inference and extracting C1' coordinates from CIF output.\n\nCode from: https://www.kaggle.com/code/zoushuxian/aido-rna-augmented-protenix","metadata":{}},{"cell_type":"code","source":"from biotite.structure.io.pdbx import CIFFile, get_structure\n\n\ndef extract_c1_atoms(cif_path):\n    \"\"\"Extract C1' atom coordinates from a CIF file.\"\"\"\n    cif_file = CIFFile.read(cif_path)\n    model = get_structure(cif_file, model=1)\n    chain = model[model.chain_id == \"A\"]\n    mask = chain.atom_name == \"C1'\"\n    c1_atoms = chain[mask]\n    df = pd.DataFrame.from_dict(c1_atoms._annot)\n    df[\"x\"] = c1_atoms.coord[:, 0]\n    df[\"y\"] = c1_atoms.coord[:, 1]\n    df[\"z\"] = c1_atoms.coord[:, 2]\n    return df[[\"res_name\", \"res_id\", \"x\", \"y\", \"z\"]]\n\n\ndef prepare_protenix_json(target_id, sequence, output_path, input_path, max_length=400):\n    \"\"\"Create Protenix input JSON for a target.\"\"\"\n    if len(sequence) <= max_length:\n        input_json = [{\n            \"sequences\": [{\n                \"rnaSequence\": {\n                    \"sequence\": sequence,\n                    \"count\": 1,\n                    \"msa\": {\n                        \"precomputed_msa_dir\": f\"{input_path}/MSA/{target_id}.MSA.fasta\",\n                        \"pairing_db\": \"rnacentral\"\n                    }\n                }\n            }],\n            \"name\": target_id,\n        }]\n    else:\n        print(f\"    Sequence too long ({len(sequence)} > {max_length}), truncating, no MSA\")\n        input_json = [{\n            \"sequences\": [{\n                \"rnaSequence\": {\n                    \"sequence\": sequence[:max_length],\n                    \"count\": 1,\n                }\n            }],\n            \"name\": target_id,\n        }]\n\n    json_path = Path(output_path) / \"input_json\" / f\"{target_id}.json\"\n    json_path.parent.mkdir(parents=True, exist_ok=True)\n    with open(json_path, \"w\") as f:\n        json.dump(input_json, f, indent=4)\n\n\ndef run_protenix_inference(target_id, sequence, output_path, input_path,\n                           seed=101, n_cycle=10, n_sample=5, n_step=200, max_length=400):\n    \"\"\"Run Protenix inference for a single target.\"\"\"\n    if IS_KAGGLE:\n        checkpoint_path = \"/kaggle/input/protenix-finetuned-rna3db-all-1599/1599_ema_0.999.pt\"\n    else:\n        checkpoint_path = f\"{DATA_PATH}/protenix_chpt/1599_ema_0.999.pt\"\n\n    output_path = Path(output_path)\n    input_json_path = output_path / \"input_json\" / f\"{target_id}.json\"\n    dump_dir = output_path / target_id\n    dump_dir.mkdir(parents=True, exist_ok=True)\n\n    use_msa = \"True\" if len(sequence) <= max_length else \"False\"\n\n    sys.argv = [\n        \"runner/inference.py\",\n        f\"--seeds={seed}\",\n        f\"--dump_dir={dump_dir}\",\n        f\"--input_json_path={input_json_path}\",\n        f\"--model.N_cycle={n_cycle}\",\n        f\"--sample_diffusion.N_sample={n_sample}\",\n        f\"--sample_diffusion.N_step={n_step}\",\n        \"--augment.use_rnalm True\",\n        f\"--use_msa {use_msa}\",\n        f\"--load_checkpoint_path={checkpoint_path}\",\n        \"\",\n    ]\n\n    from runner.inference import run\n    run()\n\n\ndef get_protenix_predictions(target_id, sequence, output_path, seed=101, n_sample=5):\n    \"\"\"Extract C1' coordinate arrays from Protenix CIF output files.\"\"\"\n    output_path = Path(output_path)\n    predictions = []\n    for i in range(n_sample):\n        cif_path = (output_path / target_id / target_id / f\"seed_{seed}\" /\n                    \"predictions\" / f\"{target_id}_seed_{seed}_sample_{i}.cif\")\n        if cif_path.exists():\n            pred_df = extract_c1_atoms(cif_path)\n            coords = np.zeros((len(sequence), 3))\n            n_atoms = min(len(pred_df), len(sequence))\n            coords[:n_atoms] = pred_df[[\"x\", \"y\", \"z\"]].values[:n_atoms]\n            predictions.append(coords)\n    return predictions\n\n\n@contextlib.contextmanager\ndef protenix_context():\n    \"\"\"Temporarily change to Protenix directory for inference.\"\"\"\n    original_dir = os.getcwd()\n    os.chdir(PROTENIX_DIR)\n    try:\n        yield\n    finally:\n        os.chdir(original_dir)\n","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Combined Prediction Pipeline\n\n**Phase 1**: Template-based predictions for all targets. Good templates\n(similarity >= 0.0 AND percent identity >= 50%) are used directly.\n\n**Phase 2**: Protenix inference for targets that need more predictions.\nOnly runs for targets where templates were insufficient.\n\n**Phase 3**: Combine template + Protenix predictions. Any remaining gaps\n(e.g. if Protenix fails) are filled with simple de novo structures.","metadata":{}},{"cell_type":"code","source":"# Template/prediction metadata tracking (global state)\ntemplate_info_dict = {}\nprediction_metadata_dict = {}\n\n\ndef record_template_info(target_id, template_id, similarity, percent_identity):\n    if target_id not in template_info_dict:\n        template_info_dict[target_id] = {\n            'template_ids': [], 'similarities': [], 'percent_identities': []\n        }\n    template_info_dict[target_id]['template_ids'].append(template_id)\n    template_info_dict[target_id]['similarities'].append(similarity)\n    template_info_dict[target_id]['percent_identities'].append(percent_identity)\n\n\ndef record_prediction_metadata(target_id, pred_num, source, template_id=None,\n                               similarity=None, percent_identity=None):\n    if target_id not in prediction_metadata_dict:\n        prediction_metadata_dict[target_id] = {}\n    prediction_metadata_dict[target_id][pred_num] = {\n        'source': source,\n        'template_id': template_id if source == 'template' else None,\n        'similarity': similarity if source == 'template' else None,\n        'percent_identity': percent_identity if source == 'template' else None,\n    }\n\n\ndef predict_with_templates(sequence, target_id, train_seqs_df, train_coords_dict,\n                           segments_map, n_predictions=5, temporal_cutoff=None):\n    \"\"\"\n    Phase 1: Generate predictions using template-based modeling.\n\n    Returns:\n        (template_predictions, n_protenix_needed, next_pred_num)\n    \"\"\"\n    predictions = []\n    pred_num = 1\n\n    print(\"\\n\" + \"-\" * 70)\n    print(f\"Target: {target_id} ({len(sequence)} nt)\")\n\n    similar_seqs = find_similar_sequences_detailed(\n        sequence, train_seqs_df, train_coords_dict,\n        temporal_cutoff=temporal_cutoff, top_n=n_predictions\n    )\n\n    if similar_seqs:\n        for i, (tmpl_id, tmpl_seq, similarity, tmpl_coords,\n                pct_id, aligned_q, aligned_t) in enumerate(similar_seqs):\n\n            if (similarity < MIN_SIMILARITY or pct_id < MIN_PERCENT_IDENTITY) and len(tmpl_seq) < 500:\n                if SHOW_ALIGNMENT_DETAILS:\n                    print(f\"  Template {i+1}: {tmpl_id} SKIPPED \"\n                          f\"(sim={similarity:.3f}, id={pct_id:.1f}%) - below threshold\")\n                break\n\n            if USE_PROTENIX and len(aligned_q) < 100 and i == 4:\n                print(f\"  Sequence length {len(aligned_q)} short, leaving 1 space for Protenix\")\n                break\n\n            record_template_info(target_id, tmpl_id, similarity, pct_id)\n            record_prediction_metadata(target_id, pred_num, 'template',\n                                       tmpl_id, similarity, pct_id)\n\n            if SHOW_ALIGNMENT_DETAILS:\n                print(f\"  Template {i+1}: {tmpl_id} (sim={similarity:.3f}, id={pct_id:.1f}%)\")\n\n            adapted = adapt_template_to_query(sequence, tmpl_seq, tmpl_coords)\n            refined = adaptive_rna_constraints(adapted, target_id, segments_map,\n                                               confidence=similarity)\n            predictions.append(refined)\n            pred_num += 1\n\n            if len(predictions) >= n_predictions:\n                break\n\n    n_from_templates = len(predictions)\n    n_needed = n_predictions - n_from_templates\n\n    if n_needed > 0:\n        print(f\"  -> {n_from_templates} from templates, {n_needed} slots for Protenix\")\n    else:\n        print(f\"  -> All {n_predictions} predictions from templates\")\n\n    return predictions, n_needed, pred_num","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def generate_predictions_batch(sequences_df, train_seqs_df, train_coords_dict,\n                               dataset_name, use_temporal_cutoff=True,\n                               protenix_output_path=None):\n    \"\"\"\n    Generate predictions combining template-based modeling with Protenix fallback.\n    \"\"\"\n    start_time = time.time()\n    total_targets = len(sequences_df)\n\n    print(f\"\\n{'='*70}\")\n    print(f\"Predicting {total_targets} {dataset_name} sequences\")\n    print(f\"{'='*70}\")\n\n    # Build segments map for targets being predicted\n    segments_map, _ = build_segments_map(sequences_df)\n\n    # ---- Phase 1: Template predictions ----\n    print(\"\\nPHASE 1: Template-based predictions\")\n\n    template_predictions = {}   # target_id -> [coords, ...]\n    protenix_queue = {}         # target_id -> (n_needed, next_pred_num, sequence)\n\n    for _, row in sequences_df.iterrows():\n        target_id = row['target_id']\n        sequence = row['sequence']\n        temporal_cutoff = row.get('temporal_cutoff', None) if use_temporal_cutoff else None\n\n        preds, n_needed, next_pred = predict_with_templates(\n            sequence, target_id, train_seqs_df, train_coords_dict,\n            segments_map, n_predictions=5, temporal_cutoff=temporal_cutoff\n        )\n\n        template_predictions[target_id] = preds\n        if n_needed > 0:\n            protenix_queue[target_id] = (n_needed, next_pred, sequence)\n\n    template_time = time.time() - start_time\n    print(f\"\\nPhase 1 done: {template_time:.1f}s | \"\n          f\"{len(protenix_queue)} targets need Protenix\")\n\n    # ---- Phase 2: Protenix predictions ----\n    protenix_predictions = {}\n\n    if protenix_queue and USE_PROTENIX:\n        print(f\"\\nPHASE 2: Protenix for {len(protenix_queue)} targets\")\n\n        if protenix_output_path is None:\n            protenix_output_path = Path(OUTPUT_PATH) / f\"{dataset_name}_protenix\"\n        protenix_output_path = Path(protenix_output_path)\n        protenix_output_path.mkdir(parents=True, exist_ok=True)\n        input_path = Path(DATA_PATH)\n\n        with protenix_context():\n            for i, (target_id, (n_needed, next_pred, sequence)) in enumerate(protenix_queue.items()):\n                print(f\"\\n  [{i+1}/{len(protenix_queue)}] {target_id} \"\n                      f\"({len(sequence)} nt, need {n_needed})\")\n\n                try:\n                    prepare_protenix_json(target_id, sequence,\n                                         protenix_output_path, input_path)\n\n                    t0 = time.time()\n                    run_protenix_inference(\n                        target_id, sequence, protenix_output_path, input_path,\n                        seed=101, n_cycle=10, n_sample=n_needed, n_step=200\n                    )\n                    print(f\"    Done in {(time.time()-t0)/60:.1f} min\")\n\n                    preds = get_protenix_predictions(\n                        target_id, sequence, protenix_output_path,\n                        seed=101, n_sample=n_needed\n                    )\n                    protenix_predictions[target_id] = preds\n                    print(f\"    Got {len(preds)} Protenix predictions\")\n\n                except Exception as e:\n                    print(f\"    Protenix FAILED: {e}\")\n                    print(f\"    Will fall back to de novo\")\n                    protenix_predictions[target_id] = None\n\n        ptx_time = time.time() - start_time - template_time\n        print(f\"\\nPhase 2 done: {ptx_time/60:.1f} min\")\n\n    elif protenix_queue and not USE_PROTENIX:\n        print(f\"\\nPHASE 2: Protenix disabled, will use de novo for \"\n              f\"{len(protenix_queue)} targets\")\n\n    # ---- Phase 3: Combine and format ----\n    print(f\"\\nPHASE 3: Combining predictions\")\n\n    all_rows = []\n\n    for _, row in sequences_df.iterrows():\n        target_id = row['target_id']\n        sequence = row['sequence']\n\n        predictions = list(template_predictions[target_id])\n        pred_num = len(predictions) + 1\n\n        # Fill with Protenix predictions\n        if target_id in protenix_queue:\n            ptx_preds = protenix_predictions.get(target_id)\n            if ptx_preds:\n                for coords in ptx_preds:\n                    record_prediction_metadata(target_id, pred_num, 'protenix')\n                    predictions.append(coords)\n                    pred_num += 1\n                    if len(predictions) >= 5:\n                        break\n\n        # Fill remaining with de novo (fallback)\n        n_denovo = 0\n        while len(predictions) < 5:\n            record_prediction_metadata(target_id, pred_num, 'de_novo')\n            seed_val = hash(target_id) % 10000 + len(predictions) * 1000\n            de_novo = generate_rna_structure(sequence, seed=seed_val)\n            refined = adaptive_rna_constraints(de_novo, target_id, segments_map,\n                                               confidence=0.2)\n            predictions.append(refined)\n            pred_num += 1\n            n_denovo += 1\n\n        if n_denovo > 0:\n            print(f\"  {target_id}: filled {n_denovo} slots with de novo\")\n\n        # Format output rows\n        for j in range(len(sequence)):\n            pred_row = {\n                'ID': f\"{target_id}_{j+1}\",\n                'resname': sequence[j],\n                'resid': j + 1,\n            }\n            for i in range(5):\n                pred_row[f'x_{i+1}'] = predictions[i][j][0]\n                pred_row[f'y_{i+1}'] = predictions[i][j][1]\n                pred_row[f'z_{i+1}'] = predictions[i][j][2]\n            all_rows.append(pred_row)\n\n    submission_df = pd.DataFrame(all_rows)\n    column_order = ['ID', 'resname', 'resid']\n    for i in range(1, 6):\n        for coord in ['x', 'y', 'z']:\n            column_order.append(f'{coord}_{i}')\n    submission_df = submission_df[column_order]\n\n    # Summary\n    total_time = time.time() - start_time\n    n_template_only = sum(1 for tid in sequences_df['target_id'] if tid not in protenix_queue)\n\n    print(f\"\\n{'='*70}\")\n    print(f\"{dataset_name.upper()} PREDICTIONS COMPLETE\")\n    print(f\"  Template-only targets: {n_template_only}\")\n    print(f\"  Targets with Protenix: {len(protenix_queue)}\")\n    print(f\"  Total residues: {len(submission_df)}\")\n    print(f\"  Runtime: {total_time:.1f}s ({total_time/60:.1f} min)\")\n    print(f\"{'='*70}\\n\")\n\n    return submission_df\n\n\nprint(\"Batch prediction function defined\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Validation Predictions\n\nPredict validation set using **only training data** as templates (no data leakage).\nProtenix fills in where templates are insufficient.","metadata":{}},{"cell_type":"code","source":"# Helper functions for evaluating validation results\n\n\"\"\"\nValidation evaluation utilities for RNA 3D structure prediction.\n\nEvaluates predictions against ground truth using TM-score, enriches results\nwith template/prediction metadata, and generates summary statistics and plots.\n\"\"\"\n\nimport os\nimport pandas as pd\n# %%\nimport os\nimport re\nimport pandas as pd\nimport shutil\nimport subprocess\n\n# https://www.kaggle.com/code/rhijudas/tm-score-permutechains\n\n# ---------------------\n# Helper: parse USalign output\n# ---------------------\ndef parse_tmscore_output(output: str):\n    \"\"\"\n    Parse USalign output to extract TM-score and RMSD.\n\n    Returns:\n        tm_score (float): The main TM-score.\n        rmsd (float or None): The RMSD value if found, else None.\n    \"\"\"\n    # Extract all TM-score lines (sometimes multiple appear in the output)\n    tm_matches = re.findall(r'TM-score=\\s+([\\d.]+)', output)\n    if len(tm_matches) < 2:\n        raise ValueError('No TM score found in USalign output')\n    tm_score = float(tm_matches[1])\n\n    # Try to extract RMSD (case-insensitive, handle variations and possible extra text)\n    rmsd_match = re.search(r'RMSD=\\s*([\\d.]+)', output, re.IGNORECASE)\n    rmsd = float(rmsd_match.group(1)) if rmsd_match else None\n\n    return tm_score, rmsd\n\n# ---------------------\n# PDB writers\n# ---------------------\n\ndef sanitize(xyz):\n    MIN_COORD=-999.999\n    MAX_COORD=9999.999\n    return min(max(xyz,MIN_COORD),MAX_COORD)\n\ndef write_target_line(atom_name, atom_serial, residue_name, chain_id, residue_num,\n                      x_coord, y_coord, z_coord, occupancy=1.0, b_factor=0.0, atom_type='P') -> str:\n    return f'ATOM  {atom_serial:>5d}  {atom_name:4s}{residue_name:>3s} {chain_id:1s}{residue_num:>4d}    {sanitize(x_coord):>8.3f}{sanitize(y_coord):>8.3f}{sanitize(z_coord):>8.3f}{occupancy:>6.2f}{b_factor:>6.2f}           {atom_type}\\n'\n    \ndef write2pdb(df: pd.DataFrame, xyz_id: int, target_path: str) -> int:\n    \"\"\"\n    Write single-chain PDB (chain 'A') using row['resid'] as residue_num.\n    Raises exceptions on invalid data.\n    \"\"\"\n    resolved_cnt = 0\n    with open(target_path, 'w') as fh:\n        for _, row in df.iterrows():\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                resolved_cnt += 1\n                resid_num = int(row['resid'])\n                fh.write(write_target_line(\"C1'\", resid_num, row['resname'], 'A', resid_num, x, y, z, atom_type='C'))\n    return resolved_cnt\n\ndef write2pdb_singlechain_native(df_native: pd.DataFrame, xyz_id: int, target_path: str) -> int:\n    \"\"\"\n    Write native single-chain using row['resid'] as residue numbers.\n    Assumes all required columns exist and are valid.\n    \"\"\"\n    df_sorted = df_native.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int').reset_index(drop=True)\n\n    resolved_cnt = 0\n    with open(target_path, 'w') as fh:\n        for _, row in df_sorted.iterrows():\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                resolved_cnt += 1\n                resid_num = int(row['resid'])\n                fh.write(write_target_line(\"C1'\", resid_num, row['resname'], 'A', resid_num, x, y, z, atom_type='C'))\n    return resolved_cnt\n\ndef write2pdb_multichain_from_solution(df_solution: pd.DataFrame, xyz_id: int, target_path: str) -> int:\n    \"\"\"\n    Write multi-chain PDB for native solution using columns 'chain' and 'copy' to assign chain letters.\n    Expects 'resid' convertible to int and chain/copy present. No fallbacks.\n    \"\"\"\n    df_sorted = df_solution.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int')\n\n    chain_map = {}\n    next_ord = ord('A')\n    written = 0\n    with open(target_path, 'w') as fh:\n        for _, row in df_sorted.iterrows():\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if not (x > -1e6 and y > -1e6 and z > -1e6):\n                continue\n            chain_val = row['chain']\n            copy_key = int(row['copy'])\n            g = (str(chain_val), copy_key)\n            if g not in chain_map:\n                if next_ord <= ord('Z'):\n                    ch = chr(next_ord)\n                else:\n                    ov = next_ord - ord('Z') - 1\n                    if ov < 26:\n                        ch = chr(ord('a') + ov)\n                    else:\n                        ch = chr(ord('0') + (ov - 26) % 10)\n                chain_map[g] = ch\n                next_ord += 1\n            chain_id = chain_map[g]\n            written += 1\n            resid_num = int(row['resid'])\n            fh.write(write_target_line(\"C1'\", resid_num, row['resname'], chain_id, resid_num, x, y, z, atom_type='C'))\n    return written\n\ndef write2pdb_multichain_from_groups(df_pred: pd.DataFrame, xyz_id: int, target_path: str, groups_list) -> (int, list):\n    \"\"\"\n    Write predicted multichain PDB based on a positional groups_list (tuple per residue: (chain, copy)).\n    Requires groups_list length == number of residues in df_pred (after sorting).\n    Returns (written_count, chain_letters_per_res).\n    \"\"\"\n    df_sorted = df_pred.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int').reset_index(drop=True)\n\n    if groups_list is None or len(groups_list) != len(df_sorted):\n        raise ValueError(\"groups_list must be provided and match number of residues in predicted df\")\n\n    chain_map = {}\n    next_ord = ord('A')\n    chain_letters = []\n    written = 0\n    with open(target_path, 'w') as fh:\n        for idx, row in df_sorted.iterrows():\n            g = groups_list[idx]\n            if isinstance(g, tuple):\n                gkey = (str(g[0]), int(g[1]))\n            else:\n                gkey = (str(g), None)\n            if gkey not in chain_map:\n                if next_ord <= ord('Z'):\n                    ch = chr(next_ord)\n                else:\n                    ov = next_ord - ord('Z') - 1\n                    if ov < 26:\n                        ch = chr(ord('a') + ov)\n                    else:\n                        ch = chr(ord('0') + (ov - 26) % 10)\n                chain_map[gkey] = ch\n                next_ord += 1\n            chain_id = chain_map[gkey]\n            chain_letters.append(chain_id)\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                written += 1\n                resid_num = int(row['resid'])\n                fh.write(write_target_line(\"C1'\", resid_num, row['resname'], chain_id, resid_num, x, y, z, atom_type='C'))\n    return written, chain_letters\n\ndef write2pdb_singlechain_permuted_pred(df_pred: pd.DataFrame, xyz_id: int, permuted_indices: list, target_path: str) -> int:\n    \"\"\"\n    Create single-chain PDB by concatenating predicted residues in permuted_indices order.\n    Output residue numbers are sequential starting at 1 and increase for every permuted position.\n    Raises exception if indices out of range.\n    \"\"\"\n    df_sorted = df_pred.copy()\n    df_sorted['__resid_int'] = df_sorted['resid'].astype(int)\n    df_sorted = df_sorted.sort_values('__resid_int').reset_index(drop=True)\n\n    written = 0\n    next_res = 1\n    with open(target_path, 'w') as fh:\n        for idx in permuted_indices:\n            if idx < 0 or idx >= len(df_sorted):\n                # strict behavior: raise error for invalid index\n                raise IndexError(f\"permuted index {idx} out of range for predicted residues\")\n            row = df_sorted.iloc[idx]\n            x = row[f'x_{xyz_id}']\n            y = row[f'y_{xyz_id}']\n            z = row[f'z_{xyz_id}']\n            out_resnum = next_res\n            if x > -1e6 and y > -1e6 and z > -1e6:\n                written += 1\n                fh.write(write_target_line(\"C1'\", out_resnum, row['resname'], 'A', out_resnum, x, y, z, atom_type='C'))\n            next_res += 1\n    return written\n\n# ---------------------\n# USalign wrappers\n# ---------------------\ndef run_usalign_raw(predicted_pdb: str, native_pdb: str, usalign_bin='USalign', align_sequence=False, tmscore=None, \n                    show_output: bool = False) -> str:\n    cmd = f'{usalign_bin} {predicted_pdb} {native_pdb} -atom \" C1\\'\"'\n    if show_output:\n        print('-'*100)\n        print(f'Running {cmd}')\n    if tmscore is not None:\n        cmd += f' -TMscore {tmscore}'\n        if int(tmscore) == 0:\n            cmd += ' -mm 1 -ter 0'\n    elif not align_sequence:\n        cmd += ' -TMscore 1'\n\n    proc = subprocess.run(cmd, shell=True, capture_output=True, text=True)\n    res = proc.stdout + proc.stderr\n    if show_output:\n        print(res)\n    return res\n\ndef parse_usalign_chain_orders(output: str):\n    \"\"\"\n    Parse USalign output for both Structure_1 and Structure_2 chain lists.\n    Returns (chain_list_structure1, chain_list_structure2).\n    Raises if parsing fails to find either line.\n    \"\"\"\n    chain1 = None\n    chain2 = None\n    for line in output.splitlines():\n        line = line.strip()\n        if line.startswith('Name of Structure_1:'):\n            parts = line.split(':')\n            clist = []\n            for part in parts[2:]:\n                token = part.strip()\n                if token == '':\n                    continue\n                token0 = token.split()[0]\n                last = token0.split(',')[-1]\n                ch = re.sub(r'[^A-Za-z0-9]', '', last)\n                if ch:\n                    clist.append(ch)\n            chain1 = clist\n        elif line.startswith('Name of Structure_2:'):\n            parts = line.split(':')\n            clist = []\n            for part in parts[2:]:\n                token = part.strip()\n                if token == '':\n                    continue\n                token0 = token.split()[0]\n                last = token0.split(',')[-1]\n                ch = re.sub(r'[^A-Za-z0-9]', '', last)\n                if ch:\n                    clist.append(ch)\n            chain2 = clist\n    if chain1 is None or chain2 is None:\n        raise ValueError(\"Failed to parse chain orders from USalign output\")\n    return chain1, chain2\n\n# ---------------------\n# Main scoring function (no try/except, no fallbacks)\n# ---------------------\ndef score(solution: pd.DataFrame, submission: pd.DataFrame, row_id_column_name: str, usalign_bin_hint: str = None, show_output: bool = False) -> float:\n    \"\"\"\n    Enhanced scoring with chain-permutation handling for multicopy targets.\n    This version contains no try/except blocks and will raise on any error.\n    \"\"\"\n    # determine usalign binary\n    if usalign_bin_hint:\n        usalign_bin = usalign_bin_hint\n    else:\n        if os.path.exists('/kaggle/input/usalign/USalign') and not os.path.exists('/kaggle/working/USalign'):\n            shutil.copy2('/kaggle/input/usalign/USalign', '/kaggle/working/USalign')\n            os.chmod('/kaggle/working/USalign', 0o755)\n        usalign_bin = '/kaggle/working/USalign' if os.path.exists('/kaggle/working/USalign') else 'USalign'\n\n    sol = solution.copy()\n    sub = submission.copy()\n    sol['target_id'] = sol['ID'].apply(lambda x: '_'.join(str(x).split('_')[:-1]))\n    sub['target_id'] = sub['ID'].apply(lambda x: '_'.join(str(x).split('_')[:-1]))\n\n    results = []\n    scores_df = pd.DataFrame()\n\n    for target_id, group_native in sol.groupby('target_id'):\n        group_predicted = sub[sub['target_id'] == target_id]\n        has_chain_copy = ('chain' in group_native.columns) and ('copy' in group_native.columns)\n        is_multicopy = has_chain_copy and (group_native['copy'].astype(float).max() > 1)\n\n        # precompute native models that have coords\n        native_with_coords = []\n        for native_cnt in range(1, 41):\n            native_pdb = f'native_{target_id}_{native_cnt}.pdb'\n            if show_output:\n                print(f\"Writing native model {native_pdb}\")\n            resolved_native = write2pdb(group_native, native_cnt, native_pdb)\n            if resolved_native > 0:\n                native_with_coords.append(native_cnt)\n            else:\n                if os.path.exists(native_pdb):\n                    os.remove(native_pdb)\n\n        if not native_with_coords:\n            raise ValueError(f\"No native models with coordinates for target {target_id}\")\n\n        best_per_pred = []\n        for pred_cnt in range(1, 6):\n            if not is_multicopy:\n                predicted_pdb = f'predicted_{target_id}_{pred_cnt}.pdb'\n                if show_output:\n                    print(f\"Writing predicted model {predicted_pdb}\")\n                resolved_pred = write2pdb(group_predicted, pred_cnt, predicted_pdb)\n                if resolved_pred <= 2:\n                    #print(f\"Predicted model {pred_cnt} for target {target_id} has insufficient coordinates\")\n                    best_per_pred.append( 0.0 )\n                    continue\n                \n                scores = []\n                for native_cnt in native_with_coords:\n                    native_pdb = f'native_{target_id}_{native_cnt}.pdb'\n                    out = run_usalign_raw(predicted_pdb, native_pdb, usalign_bin=usalign_bin, align_sequence=False, \n                                          tmscore=1, show_output=show_output)\n                    s, r = parse_tmscore_output(out)\n                    scores.append(s)\n                    idx = len(scores_df)\n                    scores_df.loc[idx, 'target_id'] = str(target_id)\n                    scores_df.loc[idx, 'native_cnt'] = str(native_cnt)\n                    scores_df.loc[idx, 'pred_cnt'] = str(pred_cnt)\n                    scores_df.loc[idx, 'tm_score'] = s\n                    scores_df.loc[idx, 'rmsd'] = r\n                best_per_pred.append(max(scores))\n\n            else:\n                # multicopy\n                # strict: require chain and copy columns convertible\n                gn_sorted = group_native.copy()\n                gn_sorted['__resid_int'] = gn_sorted['resid'].astype(int)\n                gn_sorted = gn_sorted.sort_values('__resid_int').reset_index(drop=True)\n                groups_list = []\n                for _, r in gn_sorted.iterrows():\n                    chain_val = r['chain']\n                    copy_i = int(r['copy'])\n                    groups_list.append((chain_val, copy_i))\n\n                # predicted multichain - groups_list must match predicted residue count or error\n                dfp_sorted = group_predicted.copy()\n                dfp_sorted['__resid_int'] = dfp_sorted['resid'].astype(int)\n                dfp_sorted = dfp_sorted.sort_values('__resid_int').reset_index(drop=True)\n                if len(groups_list) != len(dfp_sorted):\n                    raise ValueError(f\"groups_list length ({len(groups_list)}) does not match predicted residue count ({len(dfp_sorted)}) for target {target_id}\")\n\n                predicted_multi_pdb = f'pred_multi_{target_id}_{pred_cnt}.pdb'\n                resolved_pred_multi, pred_chain_letters = write2pdb_multichain_from_groups(group_predicted, pred_cnt, predicted_multi_pdb, groups_list)\n                if resolved_pred_multi == 0:\n                    #print(f\"Predicted multi model {pred_cnt} for target {target_id} has no coordinates\")\n                    best_per_pred.append( 0.0 )\n                    continue\n\n                scores = []\n                for native_cnt in native_with_coords:\n                    native_multi_pdb = f'native_multi_{target_id}_{native_cnt}.pdb'\n                    resolved_native_multi = write2pdb_multichain_from_solution(group_native, native_cnt, native_multi_pdb)\n                    if resolved_native_multi == 0:\n                        continue\n\n                    raw_out = run_usalign_raw(predicted_multi_pdb, native_multi_pdb, usalign_bin=usalign_bin, \n                                              align_sequence=True, tmscore=0, show_output=show_output)\n                    chain1, chain2 = parse_usalign_chain_orders(raw_out)  # will raise if parsing fails\n\n                    # build native->pred mapping chain2[i] -> chain1[i]\n                    native_to_pred = {n_ch: p_ch for n_ch, p_ch in zip(chain2, chain1)}\n\n                    # canonical native order = chain2 unique in order seen\n                    #native_chain_order = []\n                    #for ch in chain2:\n                    #    if ch not in native_chain_order:\n                    #        native_chain_order.append(ch)\n                    native_chain_order = list(native_to_pred.keys())\n                    native_chain_order.sort() # this is critical...\n\n                    # predicted chain order by following native chain A,B,...\n                    pred_chain_order = [native_to_pred[n_ch] for n_ch in native_chain_order if native_to_pred.get(n_ch) is not None]\n\n                    # construct pred_positions_by_chain\n                    pred_positions_by_chain = {}\n                    for idx, ch in enumerate(pred_chain_letters):\n                        if ch is None:\n                            continue\n                        pred_positions_by_chain.setdefault(ch, []).append(idx)\n\n                    # require that each chain in pred_chain_order exists in pred_positions_by_chain\n                    pred_chain_order = [p for p in pred_chain_order if p in pred_positions_by_chain]\n\n                    # form permuted indices by concatenation\n                    permuted_indices = []\n                    for ch in pred_chain_order:\n                        permuted_indices.extend(pred_positions_by_chain[ch])\n                    # append any remaining\n                    for idx in range(len(pred_chain_letters)):\n                        if idx not in permuted_indices:\n                            permuted_indices.append(idx)\n\n                    # write permuted single-chain predicted and native single-chain\n                    pred_single_perm = f'pred_permuted_{target_id}_{pred_cnt}_{native_cnt}.pdb'\n                    written_pred_single = write2pdb_singlechain_permuted_pred(group_predicted, pred_cnt, permuted_indices, pred_single_perm)\n                    native_single = f'native_single_{target_id}_{native_cnt}.pdb'\n                    written_native = write2pdb_singlechain_native(group_native, native_cnt, native_single)\n\n                    if written_pred_single <= 2 or written_native <= 2:\n                        raise ValueError(f\"Insufficient residues after permutation for target {target_id}, pred {pred_cnt}, native {native_cnt}\")\n\n                    out = run_usalign_raw(pred_single_perm, native_single, usalign_bin=usalign_bin, align_sequence=False, \n                                          tmscore=1, show_output=show_output)\n                    score_final, rmsd = parse_tmscore_output(out)\n                    scores.append(score_final)\n                    idx = len(scores_df)\n                    scores_df.loc[idx, 'target_id'] = str(target_id)\n                    scores_df.loc[idx, 'native_cnt'] = str(native_cnt)\n                    scores_df.loc[idx, 'pred_cnt'] = str(pred_cnt)\n                    scores_df.loc[idx, 'tm_score'] = score_final\n                    scores_df.loc[idx, 'rmsd'] = rmsd\n                best_per_pred.append(max(scores))\n\n        results.append(max(best_per_pred))\n\n    if not results:\n        pass\n        #raise ValueError(\"No targets scored\")\n\n    score_mean = float(sum(results) / len(results)) if len(results)>0 else 0.0\n    return score_mean, scores_df\n\n\n\ndef evaluate_validation(validation_labels, validation_predictions,\n                        usalign_bin,\n                        template_info_dict=None,\n                        prediction_metadata_dict=None,\n                        output_dir='.', show_output=False):\n    \"\"\"\n    Evaluate validation predictions against ground truth using TM-score.\n\n    Computes TM-scores, enriches with template/prediction metadata,\n    creates per-target summary, prints statistics, and saves CSVs.\n\n    Args:\n        validation_labels: DataFrame with ground truth labels.\n        validation_predictions: DataFrame with predicted coordinates.\n        usalign_bin: path to USalign binary.\n        template_info_dict: optional dict mapping target_id ->\n            {template_ids, similarities, percent_identities}.\n        prediction_metadata_dict: optional dict mapping target_id ->\n            {pred_num -> {source, template_id, ...}}.\n        output_dir: directory for saving CSV files (default: current dir).\n        show_output: whether to show the output of the TM-score calculation.\n        \n    Returns:\n        (mean_tm_score, validation_results, scores_df)\n        - mean_tm_score: float, competition metric\n        - validation_results: DataFrame, per-target summary\n        - scores_df: DataFrame, detailed per-prediction scores\n    \"\"\"\n    if template_info_dict is None:\n        template_info_dict = {}\n    if prediction_metadata_dict is None:\n        prediction_metadata_dict = {}\n\n    has_template_info = bool(template_info_dict)\n\n    print(\"\\n\" + \"=\" * 70)\n    print(\"EVALUATING VALIDATION PREDICTIONS USING TM-SCORE\")\n    print(\"=\" * 70 + \"\\n\")\n\n    # --- Run TM-score ---\n    mean_tm_score, scores_df = score(\n        validation_labels,\n        validation_predictions,\n        'ID',\n        usalign_bin_hint=usalign_bin,\n        show_output=show_output,\n    )\n\n    scores_df['target_id'] = scores_df['target_id'].astype(str)\n    scores_df['pred_cnt'] = scores_df['pred_cnt'].astype(int)\n    scores_df['native_cnt'] = scores_df['native_cnt'].astype(int)\n    scores_df['tm_score'] = scores_df['tm_score'].astype(float)\n    scores_df['rmsd'] = scores_df['rmsd'].astype(float)\n\n    # --- Enrich scores_df with prediction metadata ---\n    scores_df['pred_source'] = scores_df.apply(\n        lambda row: prediction_metadata_dict.get(row['target_id'], {})\n                      .get(row['pred_cnt'], {}).get('source', 'unknown'),\n        axis=1,\n    )\n    scores_df['pred_template_id'] = scores_df.apply(\n        lambda row: prediction_metadata_dict.get(row['target_id'], {})\n                      .get(row['pred_cnt'], {}).get('template_id', ''),\n        axis=1,\n    )\n    scores_df['pred_similarity'] = scores_df.apply(\n        lambda row: prediction_metadata_dict.get(row['target_id'], {})\n                      .get(row['pred_cnt'], {}).get('similarity', None),\n        axis=1,\n    )\n    scores_df['pred_percent_identity'] = scores_df.apply(\n        lambda row: prediction_metadata_dict.get(row['target_id'], {})\n                      .get(row['pred_cnt'], {}).get('percent_identity', None),\n        axis=1,\n    )\n\n    # --- Merge best-template info ---\n    template_data = []\n    for target_id in scores_df['target_id'].unique():\n        if target_id in template_info_dict:\n            info = template_info_dict[target_id]\n            if len(info['similarities']) > 0:\n                template_data.append({\n                    'target_id': target_id,\n                    'best_template': info['template_ids'][0],\n                    'alignment_score': info['similarities'][0],\n                    'percent_identity': info['percent_identities'][0],\n                })\n\n    template_df = pd.DataFrame(template_data)\n\n    if len(template_df) > 0:\n        scores_df = scores_df.merge(template_df, on='target_id', how='left')\n        scores_df['alignment_score'] = scores_df['alignment_score'].fillna(0.0)\n        scores_df['percent_identity'] = scores_df['percent_identity'].fillna(0.0)\n        scores_df['best_template'] = scores_df['best_template'].fillna('none')\n    else:\n        scores_df['alignment_score'] = 0.0\n        scores_df['percent_identity'] = 0.0\n        scores_df['best_template'] = 'none'\n\n    # --- Per-target summary ---\n    validation_results = scores_df.groupby('target_id').agg({\n        'tm_score': ['max', 'mean', 'std'],\n        'rmsd': 'min',\n        'alignment_score': 'first',\n        'percent_identity': 'first',\n        'best_template': 'first',\n    }).reset_index()\n\n    validation_results.columns = [\n        'target_id', 'best_tm_score', 'mean_tm_score', 'std_tm_score',\n        'best_rmsd', 'alignment_score', 'percent_identity', 'best_template',\n    ]\n\n    # Add sequence length\n    length_dict = {}\n    for id_val in validation_labels['ID']:\n        tid = id_val.rsplit('_', 1)[0]\n        length_dict[tid] = length_dict.get(tid, 0) + 1\n    validation_results['length'] = validation_results['target_id'].map(length_dict)\n\n    # --- Add per-prediction metadata columns ---\n    for pred_num in range(1, 6):\n        validation_results[f'pred_{pred_num}_source'] = validation_results['target_id'].apply(\n            lambda tid, pn=pred_num: prediction_metadata_dict.get(tid, {}).get(pn, {}).get('source', 'unknown')\n        )\n        validation_results[f'pred_{pred_num}_template_id'] = validation_results['target_id'].apply(\n            lambda tid, pn=pred_num: prediction_metadata_dict.get(tid, {}).get(pn, {}).get('template_id', '')\n        )\n        validation_results[f'pred_{pred_num}_similarity'] = validation_results['target_id'].apply(\n            lambda tid, pn=pred_num: prediction_metadata_dict.get(tid, {}).get(pn, {}).get('similarity', None)\n        )\n        validation_results[f'pred_{pred_num}_percent_identity'] = validation_results['target_id'].apply(\n            lambda tid, pn=pred_num: prediction_metadata_dict.get(tid, {}).get(pn, {}).get('percent_identity', None)\n        )\n\n    # --- Print statistics ---\n    _print_validation_summary(mean_tm_score, validation_results, has_template_info)\n\n    # --- Save CSVs ---\n    results_path = os.path.join(output_dir, 'validation_evaluation_results.csv')\n    details_path = os.path.join(output_dir, 'validation_detailed_scores.csv')\n    validation_results.to_csv(results_path, index=False, float_format='%.3f')\n    scores_df.to_csv(details_path, index=False, float_format='%.3f')\n    print(f\"\\n✓ Summary results saved to: {results_path}\")\n    print(f\"✓ Detailed scores saved to: {details_path}\")\n\n    return mean_tm_score, validation_results, scores_df\n\n\ndef _print_validation_summary(mean_tm_score, validation_results, has_template_info=True):\n    \"\"\"Print formatted validation summary statistics.\"\"\"\n    print(\"\\n\" + \"=\" * 70)\n    print(\"VALIDATION RESULTS SUMMARY\")\n    print(\"=\" * 70)\n    print(f\"\\nMean TM-Score (competition metric): {mean_tm_score:.4f}\")\n    print(f\"\\nTotal sequences evaluated: {len(validation_results)}\")\n\n    print(f\"\\nTM-Score Statistics (Best-of-5):\")\n    print(f\"  Mean:   {validation_results['best_tm_score'].mean():.4f}\")\n    print(f\"  Median: {validation_results['best_tm_score'].median():.4f}\")\n    print(f\"  Std:    {validation_results['best_tm_score'].std():.4f}\")\n    print(f\"  Min:    {validation_results['best_tm_score'].min():.4f}\")\n    print(f\"  Max:    {validation_results['best_tm_score'].max():.4f}\")\n\n    print(f\"\\nRMSD Statistics (Best):\")\n    print(f\"  Mean:   {validation_results['best_rmsd'].mean():.2f} Å\")\n    print(f\"  Median: {validation_results['best_rmsd'].median():.2f} Å\")\n\n    if has_template_info:\n        print(f\"\\nAlignment Score Statistics:\")\n        print(f\"  Mean:   {validation_results['alignment_score'].mean():.4f}\")\n        print(f\"  Median: {validation_results['alignment_score'].median():.4f}\")\n\n        print(f\"\\nPercent Identity Statistics:\")\n        print(f\"  Mean:   {validation_results['percent_identity'].mean():.2f}%\")\n        print(f\"  Median: {validation_results['percent_identity'].median():.2f}%\")\n\n    print(f\"\\nPerformance Breakdown:\")\n    good = (validation_results['best_tm_score'] > 0.5).sum()\n    medium = ((validation_results['best_tm_score'] >= 0.3) & (validation_results['best_tm_score'] <= 0.5)).sum()\n    poor = (validation_results['best_tm_score'] < 0.3).sum()\n    total = len(validation_results)\n\n    print(f\"  Excellent (TM > 0.5):  {good:3d} ({100 * good / total:.1f}%)\")\n    print(f\"  Good (0.3 ≤ TM ≤ 0.5): {medium:3d} ({100 * medium / total:.1f}%)\")\n    print(f\"  Poor (TM < 0.3):       {poor:3d} ({100 * poor / total:.1f}%)\")\n\n    print(f\"\\nDiversity Analysis (Ensemble Variation):\")\n    print(f\"  Mean Std across predictions: {validation_results['std_tm_score'].mean():.4f}\")\n    print(\"\\n\" + \"=\" * 70 + \"\\n\")\n\n    display_cols = ['target_id', 'length', 'best_tm_score', 'best_rmsd']\n    if has_template_info:\n        display_cols += ['alignment_score', 'percent_identity']\n    print(\"Top 10 Best Predictions:\")\n    print(validation_results.nlargest(10, 'best_tm_score')[display_cols])\n    print(\"\\nBottom 10 Worst Predictions:\")\n    print(validation_results.nsmallest(10, 'best_tm_score')[display_cols])\n\n\ndef plot_validation_results(validation_results, save_path='validation_evaluation_plots.png'):\n    \"\"\"\n    Generate visualization of validation results.\n\n    Produces a 6-panel figure when template info is available (alignment_score\n    and percent_identity columns have non-zero values), or a 4-panel figure\n    otherwise.\n\n    Args:\n        validation_results: DataFrame from evaluate_validation.\n        save_path: file path to save the figure.\n    \"\"\"\n    import matplotlib.pyplot as plt\n    import seaborn as sns\n\n    sns.set_style('whitegrid')\n\n    has_template_info = validation_results['alignment_score'].any()\n\n    if has_template_info:\n        fig, axes = plt.subplots(2, 3, figsize=(20, 12))\n    else:\n        fig, axes = plt.subplots(2, 2, figsize=(16, 12))\n\n    # 1. TM-Score Distribution\n    ax = axes[0, 0]\n    ax.hist(validation_results['best_tm_score'], bins=30, edgecolor='black', alpha=0.7)\n    ax.axvline(0.5, color='red', linestyle='--', linewidth=2, label='TM=0.5 (same fold threshold)')\n    ax.axvline(validation_results['best_tm_score'].mean(), color='green', linestyle='--', linewidth=2,\n               label=f'Mean={validation_results[\"best_tm_score\"].mean():.3f}')\n    ax.set_xlabel('TM-Score (Best-of-5)', fontsize=12)\n    ax.set_ylabel('Count', fontsize=12)\n    ax.set_title('Distribution of TM-Scores on Validation Set', fontsize=14, fontweight='bold')\n    ax.legend()\n    ax.grid(True, alpha=0.3)\n\n    # 2. TM-Score vs Sequence Length\n    ax = axes[0, 1]\n    scatter = ax.scatter(validation_results['length'], validation_results['best_tm_score'],\n                         c=validation_results['best_tm_score'], cmap='RdYlGn', s=50, alpha=0.6)\n    ax.axhline(0.5, color='red', linestyle='--', linewidth=1, alpha=0.5)\n    ax.set_xlabel('Sequence Length (nucleotides)', fontsize=12)\n    ax.set_ylabel('TM-Score (Best-of-5)', fontsize=12)\n    ax.set_title('TM-Score vs Sequence Length', fontsize=14, fontweight='bold')\n    plt.colorbar(scatter, ax=ax, label='TM-Score')\n    ax.grid(True, alpha=0.3)\n\n    if has_template_info:\n        # 3. TM-Score vs Alignment Score\n        ax = axes[0, 2]\n        scatter = ax.scatter(validation_results['alignment_score'], validation_results['best_tm_score'],\n                             c=validation_results['percent_identity'], cmap='viridis', s=50, alpha=0.6)\n        ax.set_xlabel('Alignment Score', fontsize=12)\n        ax.set_ylabel('TM-Score (Best-of-5)', fontsize=12)\n        ax.set_title('TM-Score vs Template Alignment Score', fontsize=14, fontweight='bold')\n        plt.colorbar(scatter, ax=ax, label='Percent Identity')\n        ax.grid(True, alpha=0.3)\n\n    # 4. RMSD Distribution\n    ax = axes[1, 0]\n    rmsd_filtered = validation_results[validation_results['best_rmsd'] < 100]\n    ax.hist(rmsd_filtered['best_rmsd'], bins=30, edgecolor='black', alpha=0.7, color='orange')\n    ax.axvline(rmsd_filtered['best_rmsd'].mean(), color='red', linestyle='--', linewidth=2,\n               label=f'Mean={rmsd_filtered[\"best_rmsd\"].mean():.2f}Å')\n    ax.set_xlabel('RMSD (Å)', fontsize=12)\n    ax.set_ylabel('Count', fontsize=12)\n    ax.set_title('Distribution of RMSD Values (<100Å)', fontsize=14, fontweight='bold')\n    ax.legend()\n    ax.grid(True, alpha=0.3)\n\n    if has_template_info:\n        # 5. Percent Identity Distribution\n        ax = axes[1, 1]\n        ax.hist(validation_results['percent_identity'], bins=30, edgecolor='black', alpha=0.7, color='purple')\n        ax.axvline(validation_results['percent_identity'].mean(), color='red', linestyle='--', linewidth=2,\n                   label=f'Mean={validation_results[\"percent_identity\"].mean():.1f}%')\n        ax.set_xlabel('Percent Identity (%)', fontsize=12)\n        ax.set_ylabel('Count', fontsize=12)\n        ax.set_title('Distribution of Template Percent Identity', fontsize=14, fontweight='bold')\n        ax.legend()\n        ax.grid(True, alpha=0.3)\n\n        # 6. TM-Score vs Percent Identity\n        ax = axes[1, 2]\n        scatter = ax.scatter(validation_results['percent_identity'], validation_results['best_tm_score'],\n                             c=validation_results['alignment_score'], cmap='plasma', s=50, alpha=0.6)\n        ax.set_xlabel('Percent Identity (%)', fontsize=12)\n        ax.set_ylabel('TM-Score (Best-of-5)', fontsize=12)\n        ax.set_title('TM-Score vs Template Percent Identity', fontsize=14, fontweight='bold')\n        plt.colorbar(scatter, ax=ax, label='Alignment Score')\n        ax.grid(True, alpha=0.3)\n    else:\n        # 5 (alt). TM-Score vs RMSD\n        ax = axes[1, 1]\n        scatter = ax.scatter(validation_results['best_rmsd'], validation_results['best_tm_score'],\n                             c=validation_results['length'], cmap='viridis', s=50, alpha=0.6)\n        ax.set_xlabel('RMSD (Å)', fontsize=12)\n        ax.set_ylabel('TM-Score (Best-of-5)', fontsize=12)\n        ax.set_title('TM-Score vs RMSD', fontsize=14, fontweight='bold')\n        plt.colorbar(scatter, ax=ax, label='Sequence Length')\n        ax.grid(True, alpha=0.3)\n\n    plt.tight_layout()\n    plt.savefig(save_path, dpi=150, bbox_inches='tight')\n    plt.show()\n\n    print(f\"✓ Visualization saved to: {save_path}\")","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if SHOW_VALIDATION:\n    # Reset metadata for clean validation run\n    template_info_dict.clear()\n    prediction_metadata_dict.clear()\n    if DEBUG:\n        print(validation_seqs.sequence.apply(len).sort_values())\n        shortest_ids = (validation_seqs\n                        .assign(seq_len=validation_seqs['sequence'].str.len())\n                        .nsmallest(5, 'seq_len')['target_id'])\n        validation_seqs = validation_seqs[validation_seqs['target_id'].isin(shortest_ids)].reset_index(drop=True)\n        validation_labels = validation_labels[validation_labels['ID'].str.rsplit('_', n=1).str[0].isin(shortest_ids)].reset_index(drop=True)\n        print(f\"DEBUG: Using {len(validation_seqs)} shortest sequences for validation\")\n\n\n    validation_predictions = generate_predictions_batch(\n        validation_seqs,\n        train_seqs,\n        train_coords_dict,\n        dataset_name=\"validation\",\n        use_temporal_cutoff=True,\n    )\n\n    validation_predictions.to_csv('validation_predictions.csv', index=False)\n    print(\"Saved: validation_predictions.csv\")\n    validation_predictions.head(10)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if SHOW_VALIDATION:\n    \n    mean_tm_score, validation_results, scores_df = evaluate_validation(\n        validation_labels,\n        validation_predictions,\n        usalign_bin=USALIGN_BIN,\n        template_info_dict=template_info_dict,\n        prediction_metadata_dict=prediction_metadata_dict,\n    )\n\n    print(\"Prediction metadata sample:\")\n    metadata_cols = ['target_id'] + [f'pred_{i}_source' for i in range(1, 6)]\n    print(validation_results[metadata_cols].head(10))","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"code","source":"if SHOW_VALIDATION:\n    plot_validation_results(validation_results)","metadata":{},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Test Predictions (Competition Submission)\n\nPredict test set using **combined training + validation data** as templates.\nProtenix fills remaining prediction slots.","metadata":{}},{"cell_type":"code","source":"if MAKE_SUBMISSION:\n    # Reset metadata for test run\n    template_info_dict.clear()\n    prediction_metadata_dict.clear()\n\n    test_predictions = generate_predictions_batch(\n        test_seqs,\n        combined_seqs,\n        combined_coords_dict,\n        dataset_name=\"test\",\n        use_temporal_cutoff=False,\n    )\n\n    test_predictions.to_csv('submission.csv', index=False)\n    print(\"Saved: submission.csv\")\n    test_predictions.head(10)","metadata":{},"outputs":[],"execution_count":null}]}