{"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":[{"sourceId":118765,"databundleVersionId":15231210,"sourceType":"competition"}],"dockerImageVersionId":31234,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# # CPU Monte Carlo Baseline v2 ( v1 ~0.127)","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, math\nimport numpy as np\nimport pandas as pd\n\n# -----------------------------\n# IO\n# -----------------------------\nPAIR_SET = {(\"A\",\"U\"),(\"U\",\"A\"),(\"G\",\"C\"),(\"C\",\"G\"),(\"G\",\"U\"),(\"U\",\"G\")}\n\ndef read_test_sequences(path=\"test_sequences.csv\"):\n    df = pd.read_csv(path)\n    cols = [c.lower() for c in df.columns]\n    id_col = df.columns[cols.index(\"id\")] if \"id\" in cols else df.columns[0]\n    seq_col = df.columns[cols.index(\"sequence\")] if \"sequence\" in cols else df.columns[1]\n    out = df[[id_col, seq_col]].rename(columns={id_col: \"id\", seq_col: \"sequence\"})\n    out[\"sequence\"] = out[\"sequence\"].astype(str).str.upper().str.strip()\n    return out\n\ndef can_pair(a, b):\n    return (a, b) in PAIR_SET\n\ndef is_gc(a,b): return (a,b) in {(\"G\",\"C\"),(\"C\",\"G\")}\ndef is_au(a,b): return (a,b) in {(\"A\",\"U\"),(\"U\",\"A\")}\ndef is_gu(a,b): return (a,b) in {(\"G\",\"U\"),(\"U\",\"G\")}\n\n# -----------------------------\n# 1) Secondary structure via DP (Nussinov + stacking reward)\n# -----------------------------\ndef nussinov_stacking(seq, min_loop=3):\n    \"\"\"\n    Noncrossing maximum-score pairing.\n    Score = basepair_score + stacking_bonus if (i+1,j-1) also paired.\n    \"\"\"\n    n = len(seq)\n    dp = np.zeros((n, n), dtype=np.float32)\n    bt = np.zeros((n, n), dtype=np.int8)  # 0 i unpaired, 1 j unpaired, 2 pair, 3 split\n\n    def bp_score(a, b):\n        if not can_pair(a, b): return -1e9\n        if is_gc(a,b): return 2.4\n        if is_au(a,b): return 2.0\n        if is_gu(a,b): return 1.4  # weaker\n        return 1.6\n\n    for L in range(1, n):\n        for i in range(0, n - L):\n            j = i + L\n            best = dp[i+1, j]  # i unpaired\n            dec = 0\n\n            v = dp[i, j-1]     # j unpaired\n            if v > best:\n                best = v; dec = 1\n\n            # i-j paired\n            if j - i > min_loop and can_pair(seq[i], seq[j]):\n                base = bp_score(seq[i], seq[j])\n                inside = dp[i+1, j-1] if (i+1 <= j-1) else 0.0\n\n                # stacking bonus encourages helices\n                stack = 0.0\n                if i+1 < j-1 and can_pair(seq[i+1], seq[j-1]):\n                    stack = 0.6\n\n                v = inside + base + stack\n                if v > best:\n                    best = v; dec = 2\n\n            # split\n            # (for speed, you can restrict k range for huge n, but test has manageable count)\n            best_split = best\n            best_k = -1\n            for k in range(i+1, j):\n                v = dp[i, k] + dp[k+1, j]\n                if v > best_split:\n                    best_split = v\n                    best_k = k\n            if best_k != -1:\n                best = best_split\n                dec = 3\n\n            dp[i, j] = best\n            bt[i, j] = dec\n\n    pairs = {}\n\n    def traceback(i, j):\n        if i >= j: return\n        dec = bt[i, j]\n        if dec == 0:\n            traceback(i+1, j)\n        elif dec == 1:\n            traceback(i, j-1)\n        elif dec == 2:\n            pairs[i] = j; pairs[j] = i\n            traceback(i+1, j-1)\n        else:\n            # find best k again\n            best = -1e9\n            best_k = None\n            for k in range(i+1, j):\n                v = dp[i, k] + dp[k+1, j]\n                if v > best:\n                    best = v; best_k = k\n            traceback(i, best_k)\n            traceback(best_k+1, j)\n\n    traceback(0, n-1)\n    # Return list (i<j)\n    out = sorted([(i,j) for i,j in pairs.items() if i < j], key=lambda x: x[0])\n    return out\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T09:06:30.289119Z","iopub.execute_input":"2026-01-22T09:06:30.289404Z","iopub.status.idle":"2026-01-22T09:06:31.442773Z","shell.execute_reply.started":"2026-01-22T09:06:30.289381Z","shell.execute_reply":"2026-01-22T09:06:31.441997Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 2) Helix-first initializer (A-form-ish C1' trace)\n# -----------------------------\ndef build_helix_init(seq, pairs, rng, bond_target=5.9):\n    \"\"\"\n    Place paired regions on a loose helix, connect loops by interpolation.\n    This gives MC a much better starting point than a random walk.\n    \"\"\"\n    n = len(seq)\n    coords = np.full((n,3), np.nan, dtype=np.float32)\n\n    # Group pairs into stems: consecutive (i,i+1...) with (j,j-1...)\n    pairs = sorted(pairs)\n    stems = []\n    k = 0\n    while k < len(pairs):\n        i,j = pairs[k]\n        stem_i = [i]; stem_j = [j]\n        k2 = k+1\n        while k2 < len(pairs):\n            i2,j2 = pairs[k2]\n            if i2 == stem_i[-1] + 1 and j2 == stem_j[-1] - 1:\n                stem_i.append(i2); stem_j.append(j2)\n                k2 += 1\n            else:\n                break\n        stems.append((stem_i, stem_j))\n        k = k2\n\n    # Helix parameters (jitter per structure)\n    rise = 2.8 * (1.0 + rng.normal(0, 0.05))\n    twist = np.deg2rad(32.7 * (1.0 + rng.normal(0, 0.06)))\n    radius = 8.0 * (1.0 + rng.normal(0, 0.07))\n\n    # Place each stem in its own local frame along z, later we pack them in space.\n    stem_blocks = []\n    for s_idx, (si, sj) in enumerate(stems):\n        L = len(si)\n        # local coords for one strand along helix, paired strand opposite phase\n        block = {}\n        phase0 = rng.uniform(0, 2*np.pi)\n        z0 = 0.0\n        for t in range(L):\n            ang = phase0 + t*twist\n            z = z0 + t*rise\n            p1 = np.array([radius*np.cos(ang), radius*np.sin(ang), z], dtype=np.float32)\n            p2 = np.array([radius*np.cos(ang+np.pi), radius*np.sin(ang+np.pi), z], dtype=np.float32)\n            block[si[t]] = p1\n            block[sj[t]] = p2\n        stem_blocks.append(block)\n\n    # Pack stems in 3D: place each block with random rotation + translation, avoid overlap crudely\n    placed = np.zeros((0,3), dtype=np.float32)\n    for block in stem_blocks:\n        idxs = sorted(block.keys())\n        pts = np.stack([block[i] for i in idxs], axis=0)\n\n        # random rotation\n        axis = rng.normal(size=3).astype(np.float32)\n        axis /= (np.linalg.norm(axis)+1e-9)\n        angle = rng.uniform(0, 2*np.pi)\n        R = rodrigues_R(axis, angle)\n\n        pts = (pts @ R.T)\n\n        # translation: spread stems apart\n        tvec = rng.normal(size=3).astype(np.float32)\n        tvec /= (np.linalg.norm(tvec)+1e-9)\n        tvec *= rng.uniform(15.0, 35.0)\n        pts = pts + tvec\n\n        # assign\n        for ii,p in zip(idxs, pts):\n            coords[ii] = p\n\n    # Fill missing residues (loops/unpaired) by smooth interpolation between nearest known anchors\n    known = np.where(~np.isnan(coords[:,0]))[0].tolist()\n    if not known:\n        # fallback: random walk\n        coords = init_random_walk(n, bond_target, rng)\n        return coords\n\n    # Ensure ends have anchors by setting them if missing\n    if np.isnan(coords[0,0]):\n        coords[0] = coords[known[0]] + rng.normal(scale=2.0, size=3).astype(np.float32)\n        known = sorted(set(known + [0]))\n    if np.isnan(coords[-1,0]):\n        coords[-1] = coords[known[-1]] + rng.normal(scale=2.0, size=3).astype(np.float32)\n        known = sorted(set(known + [n-1]))\n\n    known = sorted(known)\n    # linear interpolation then small noise, then enforce approximate bond lengths by local smoothing\n    for a, b in zip(known[:-1], known[1:]):\n        if b == a+1: \n            continue\n        pa = coords[a].copy()\n        pb = coords[b].copy()\n        m = b - a\n        for t in range(1, m):\n            u = t / m\n            coords[a+t] = (1-u)*pa + u*pb + rng.normal(scale=1.2, size=3).astype(np.float32)\n\n    coords -= coords.mean(axis=0, keepdims=True)\n\n    # Quick bond-length smoothing: pull each step toward bond_target\n    for _ in range(2):\n        dif = coords[1:] - coords[:-1]\n        d = np.linalg.norm(dif, axis=1, keepdims=True) + 1e-9\n        dif_unit = dif / d\n        coords[1:] = coords[:-1] + dif_unit * bond_target\n\n    coords -= coords.mean(axis=0, keepdims=True)\n    return coords.astype(np.float32)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T09:06:35.695994Z","iopub.execute_input":"2026-01-22T09:06:35.6963Z","iopub.status.idle":"2026-01-22T09:06:35.713035Z","shell.execute_reply.started":"2026-01-22T09:06:35.69627Z","shell.execute_reply":"2026-01-22T09:06:35.712326Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 3) Moves (yours + one helix-aware move)\n# -----------------------------\ndef rodrigues_R(axis, angle):\n    axis = axis / (np.linalg.norm(axis) + 1e-9)\n    ax = axis\n    K = np.array([[0, -ax[2], ax[1]],\n                  [ax[2], 0, -ax[0]],\n                  [-ax[1], ax[0], 0]], dtype=np.float32)\n    I = np.eye(3, dtype=np.float32)\n    return I + math.sin(angle) * K + (1 - math.cos(angle)) * (K @ K)\n\ndef propose_local(coords, rng, sigma):\n    prop = coords.copy()\n    i = int(rng.randint(0, coords.shape[0]))\n    prop[i] += rng.normal(scale=sigma, size=3).astype(np.float32)\n    return prop\n\ndef propose_segment(coords, rng, sigma, max_len=25):\n    N = coords.shape[0]\n    i = int(rng.randint(0, N))\n    j = int(rng.randint(0, N))\n    if i > j: i, j = j, i\n    if j - i > max_len:\n        j = i + max_len\n    prop = coords.copy()\n    delta = rng.normal(scale=sigma, size=3).astype(np.float32)\n    prop[i:j+1] += delta\n    return prop\n\ndef propose_pivot(coords, rng, sigma_angle=0.2):\n    N = coords.shape[0]\n    pivot = int(rng.randint(0, N-1))\n    prop = coords.copy()\n    axis = rng.normal(size=3).astype(np.float32)\n    angle = float(rng.normal(scale=sigma_angle))\n    R = rodrigues_R(axis, angle)\n    origin = prop[pivot].copy()\n    tail = prop[pivot+1:] - origin\n    prop[pivot+1:] = (tail @ R.T) + origin\n    return prop\n\ndef propose_crankshaft(coords, rng, sigma_angle=0.35, min_gap=3, max_gap=20):\n    N = coords.shape[0]\n    if N < (min_gap + 3):\n        return coords.copy()\n    i = int(rng.randint(0, N - min_gap - 1))\n    gap = int(rng.randint(min_gap, max_gap + 1))\n    j = min(N - 1, i + gap)\n    if j - i < min_gap:\n        return coords.copy()\n\n    prop = coords.copy()\n    a = prop[i].copy()\n    b = prop[j].copy()\n    axis = (b - a).astype(np.float32)\n    axis /= (np.linalg.norm(axis) + 1e-9)\n    angle = float(rng.normal(scale=sigma_angle))\n    R = rodrigues_R(axis, angle)\n\n    seg = prop[i+1:j] - a\n    prop[i+1:j] = (seg @ R.T) + a\n    return prop\n\ndef propose_helix_rigidbody(coords, rng, pairs, max_block_len=50):\n    \"\"\"\n    Pick a paired residue i and rotate/translate a contiguous neighborhood.\n    Helps move whole helices without breaking local geometry too much.\n    \"\"\"\n    if not pairs:\n        return coords.copy()\n    N = coords.shape[0]\n    prop = coords.copy()\n\n    # pick a random pair and define a block around it\n    i, j = pairs[int(rng.randint(0, len(pairs)))]\n    center = (i + j)//2\n    L = int(rng.randint(10, max_block_len+1))\n    a = max(0, center - L//2)\n    b = min(N-1, center + L//2)\n\n    block = prop[a:b+1]\n    origin = block.mean(axis=0)\n\n    axis = rng.normal(size=3).astype(np.float32)\n    axis /= (np.linalg.norm(axis)+1e-9)\n    angle = rng.normal(scale=0.25)\n    R = rodrigues_R(axis, float(angle))\n\n    block2 = (block - origin) @ R.T + origin\n    block2 += rng.normal(scale=0.4, size=3).astype(np.float32)  # small translation\n    prop[a:b+1] = block2\n    return prop\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T09:07:01.177588Z","iopub.execute_input":"2026-01-22T09:07:01.177908Z","iopub.status.idle":"2026-01-22T09:07:01.192201Z","shell.execute_reply.started":"2026-01-22T09:07:01.17788Z","shell.execute_reply":"2026-01-22T09:07:01.191441Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 4) Energy (less noise + stacking + paired stiffness)\n# -----------------------------\ndef energy(coords: np.ndarray,\n           pairs,\n           rng,\n           bond_target=5.9, bond_k=10.0,\n           angle_k=3.0, target_cos=0.65,\n           pair_k=1.4, pair_lo=8.0, pair_hi=13.0,\n           stack_k=0.8,     # NEW: encourage stacked pairs to keep similar distances\n           paired_smooth_k=0.8,  # NEW: keep paired regions smoother\n           rep_cutoff=3.0, rep_k=0.12,\n           rep_local_window=60  # NEW: vectorized local repulsion within a window\n           ):\n    N = coords.shape[0]\n    E = 0.0\n\n    # (a) bond length\n    dif = coords[1:] - coords[:-1]\n    d = np.linalg.norm(dif, axis=1) + 1e-9\n    E += bond_k * float(np.sum((d - bond_target) ** 2))\n\n    # (b) smoothness\n    vn = dif / (np.linalg.norm(dif, axis=1, keepdims=True) + 1e-9)\n    cosang = np.sum(vn[:-1] * vn[1:], axis=1)\n    E += angle_k * float(np.sum((cosang - target_cos) ** 2))\n\n    # (c) base-pair distance (flat-bottom)\n    if pairs:\n        ii = np.fromiter((i for i, _ in pairs), dtype=np.int32)\n        jj = np.fromiter((j for _, j in pairs), dtype=np.int32)\n        rij = coords[ii] - coords[jj]\n        dij = np.linalg.norm(rij, axis=1) + 1e-9\n        below = np.clip(pair_lo - dij, 0, None)\n        above = np.clip(dij - pair_hi, 0, None)\n        E += pair_k * float(np.sum(below**2 + above**2))\n\n        # (d) stacking: for consecutive stacked pairs (i,i+1) with (j,j-1)\n        # penalize big change in pair distance between stacked layers -> stabilizes helices\n        # Build a quick map for lookup\n        pair_map = {}\n        for i,j in pairs:\n            pair_map[i]=j; pair_map[j]=i\n        stack_terms = 0.0\n        cnt = 0\n        for i,j in pairs:\n            # Only count one direction i<j\n            if i > j: \n                continue\n            i2 = i+1\n            j2 = j-1\n            if i2 < j2 and (i2 in pair_map) and (pair_map[i2] == j2):\n                d1 = np.linalg.norm(coords[i] - coords[j]) + 1e-9\n                d2 = np.linalg.norm(coords[i2] - coords[j2]) + 1e-9\n                stack_terms += (d1 - d2)**2\n                cnt += 1\n        if cnt > 0:\n            E += stack_k * float(stack_terms)\n\n        # (e) paired-region smoothness: if i is paired, encourage its backbone tangent to be smoother\n        # (cheap: extra weight on angles near paired residues)\n        paired = np.zeros(N, dtype=np.float32)\n        for i,j in pairs:\n            paired[i] = 1.0; paired[j] = 1.0\n        # angles are at residue 1..N-2; weight by whether residue is paired\n        w = 0.5 + paired[1:-1]\n        E += paired_smooth_k * float(np.sum(w * (cosang - target_cos)**2))\n\n    # (f) local repulsion (vectorized within a sequence-distance window)\n    # This avoids high-variance sampled repulsion and is still cheap.\n    W = rep_local_window\n    for i in range(0, N):\n        j0 = i + 3\n        j1 = min(N, i + W)\n        if j0 >= j1:\n            continue\n        rij = coords[i] - coords[j0:j1]\n        dij = np.linalg.norm(rij, axis=1) + 1e-9\n        mask = dij < rep_cutoff\n        if np.any(mask):\n            x = (rep_cutoff - dij[mask])\n            E += rep_k * float(np.sum(x*x))\n\n    return float(E)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T09:07:09.100928Z","iopub.execute_input":"2026-01-22T09:07:09.10144Z","iopub.status.idle":"2026-01-22T09:07:09.113086Z","shell.execute_reply.started":"2026-01-22T09:07:09.101412Z","shell.execute_reply":"2026-01-22T09:07:09.112451Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 5) Init + MC Anneal + Nested MC (depth=2)\n# -----------------------------\ndef init_random_walk(N, bond_target, rng):\n    coords = np.zeros((N, 3), dtype=np.float32)\n    for i in range(1, N):\n        v = rng.normal(size=3).astype(np.float32)\n        v /= (np.linalg.norm(v) + 1e-9)\n        coords[i] = coords[i-1] + bond_target * v\n    coords -= coords.mean(axis=0, keepdims=True)\n    return coords\n\ndef one_mc_step(coords, pairs, rng, T, sigma, sigma_angle, move_probs):\n    u = rng.rand()\n    if u < move_probs[0]:\n        prop = propose_local(coords, rng, sigma)\n    elif u < move_probs[1]:\n        prop = propose_crankshaft(coords, rng, sigma_angle=sigma_angle)\n    elif u < move_probs[2]:\n        prop = propose_segment(coords, rng, sigma * 0.6)\n    elif u < move_probs[3]:\n        prop = propose_pivot(coords, rng, sigma_angle=sigma_angle)\n    else:\n        prop = propose_helix_rigidbody(coords, rng, pairs)\n\n    return prop\n\ndef run_anneal(coords, pairs, rng, steps,\n               T_start, T_end,\n               sigma_start, sigma_end,\n               angle_k, pair_k,\n               use_nested=False,\n               nested_rollout=24):\n    # Move CDF: local, crank, segment, pivot, helix-rigidbody\n    move_cdf = [0.42, 0.68, 0.84, 0.94, 1.00]\n\n    E = energy(coords, pairs, rng, angle_k=angle_k, pair_k=pair_k)\n    best = coords.copy()\n    best_E = E\n\n    for t in range(steps):\n        frac = t / max(1, steps - 1)\n        T = T_start * (T_end / T_start) ** frac\n        sigma = sigma_start * (sigma_end / sigma_start) ** frac\n        sigma_angle = 0.35 * (1 - frac) + 0.08\n\n        if use_nested:\n            # Nested MC depth-2: choose among K candidates using short rollout\n            K = 5\n            best_prop = None\n            best_roll_E = 1e18\n            # (Important) use a deterministic fork of RNG to keep reproducible\n            base_state = rng.get_state()\n            for k in range(K):\n                rng.set_state(base_state)\n                # jump RNG deterministically per k\n                rng.randint(0, 10**9, size=(k+1,))\n                prop0 = one_mc_step(coords, pairs, rng, T, sigma, sigma_angle, move_cdf)\n\n                # rollout from prop0 with a few greedy-ish steps at slightly lower T\n                c = prop0\n                e = energy(c, pairs, rng, angle_k=angle_k, pair_k=pair_k)\n                T2 = max(0.15, 0.7*T)\n                for _ in range(nested_rollout):\n                    p = one_mc_step(c, pairs, rng, T2, sigma*0.8, sigma_angle*0.9, move_cdf)\n                    e2 = energy(p, pairs, rng, angle_k=angle_k, pair_k=pair_k)\n                    dE = e2 - e\n                    if dE <= 0 or rng.rand() < math.exp(-dE / max(1e-9, T2)):\n                        c, e = p, e2\n                if e < best_roll_E:\n                    best_roll_E = e\n                    best_prop = prop0\n\n            rng.set_state(base_state)\n            prop = best_prop\n        else:\n            prop = one_mc_step(coords, pairs, rng, T, sigma, sigma_angle, move_cdf)\n\n        E_new = energy(prop, pairs, rng, angle_k=angle_k, pair_k=pair_k)\n        dE = E_new - E\n        if dE <= 0 or rng.rand() < math.exp(-dE / max(1e-9, T)):\n            coords = prop\n            E = E_new\n            if E < best_E:\n                best_E = E\n                best = coords.copy()\n\n        if (t + 1) % 400 == 0:\n            coords -= coords.mean(axis=0, keepdims=True)\n\n    best -= best.mean(axis=0, keepdims=True)\n    return best, best_E\n\ndef budget_by_length(N):\n    # slightly stronger budgets for small/medium (where you can win most on LB)\n    if N <= 120:\n        return dict(\n            max_span=280, bond_target=5.9,\n            steps1=1800, steps2=2200,\n            T1s=3.8, T1e=0.9,  sig1s=0.85, sig1e=0.30, angle1=2.2, pair1=1.0,\n            T2s=1.0, T2e=0.22, sig2s=0.30, sig2e=0.12, angle2=3.2, pair2=1.5,\n            use_nested=True, rollout=22\n        )\n    if N <= 300:\n        return dict(\n            max_span=260, bond_target=5.9,\n            steps1=1100, steps2=1400,\n            T1s=3.6, T1e=1.0,  sig1s=0.80, sig1e=0.28, angle1=2.0, pair1=1.0,\n            T2s=1.0, T2e=0.26, sig2s=0.28, sig2e=0.12, angle2=3.0, pair2=1.4,\n            use_nested=True, rollout=18\n        )\n    if N <= 700:\n        return dict(\n            max_span=220, bond_target=5.9,\n            steps1=520, steps2=650,\n            T1s=3.2, T1e=1.2,  sig1s=0.72, sig1e=0.24, angle1=1.7, pair1=0.9,\n            T2s=0.95, T2e=0.30, sig2s=0.24, sig2e=0.11, angle2=2.4, pair2=1.2,\n            use_nested=False, rollout=0\n        )\n    if N <= 1800:\n        return dict(\n            max_span=180, bond_target=5.9,\n            steps1=190, steps2=240,\n            T1s=2.8, T1e=1.4,  sig1s=0.55, sig1e=0.20, angle1=1.2, pair1=0.7,\n            T2s=0.85, T2e=0.38, sig2s=0.20, sig2e=0.10, angle2=1.8, pair2=0.95,\n            use_nested=False, rollout=0\n        )\n    return dict(\n        max_span=140, bond_target=5.9,\n        steps1=70, steps2=90,\n        T1s=2.4, T1e=1.5,  sig1s=0.45, sig1e=0.18, angle1=1.0, pair1=0.55,\n        T2s=0.8, T2e=0.45, sig2s=0.18, sig2e=0.09, angle2=1.4, pair2=0.75,\n        use_nested=False, rollout=0\n    )\n\ndef mc_two_stage(seq, pairs, seed, budget):\n    rng = np.random.RandomState(seed)\n    N = len(seq)\n\n    # Better init: helix-first based on pairs\n    coords = build_helix_init(seq, pairs, rng, bond_target=budget[\"bond_target\"])\n\n    # Stage 1: explore\n    coords, _ = run_anneal(\n        coords, pairs, rng,\n        steps=budget[\"steps1\"],\n        T_start=budget[\"T1s\"], T_end=budget[\"T1e\"],\n        sigma_start=budget[\"sig1s\"], sigma_end=budget[\"sig1e\"],\n        angle_k=budget[\"angle1\"],\n        pair_k=budget[\"pair1\"],\n        use_nested=budget.get(\"use_nested\", False),\n        nested_rollout=budget.get(\"rollout\", 0)\n    )\n\n    # Stage 2: refine (no nested; tighter)\n    coords, E2 = run_anneal(\n        coords, pairs, rng,\n        steps=budget[\"steps2\"],\n        T_start=budget[\"T2s\"], T_end=budget[\"T2e\"],\n        sigma_start=budget[\"sig2s\"], sigma_end=budget[\"sig2e\"],\n        angle_k=budget[\"angle2\"],\n        pair_k=budget[\"pair2\"],\n        use_nested=False,\n        nested_rollout=0\n    )\n    return coords.astype(np.float32), E2\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T09:07:19.555627Z","iopub.execute_input":"2026-01-22T09:07:19.555942Z","iopub.status.idle":"2026-01-22T09:07:19.576977Z","shell.execute_reply.started":"2026-01-22T09:07:19.555914Z","shell.execute_reply":"2026-01-22T09:07:19.576147Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 6) Diverse 5-structure strategy\n# -----------------------------\ndef make_submission(test_df, n_structures=5, base_seed=2026):\n    rows = []\n    for r, (sid, seq) in enumerate(zip(test_df[\"id\"], test_df[\"sequence\"])):\n        N = len(seq)\n        bud = budget_by_length(N)\n\n        # Base pairing from DP\n        pairs0 = nussinov_stacking(seq, min_loop=3)\n\n        # Create 5 diverse pair-sets: tweak min_loop and (for some) drop weakest GU pairs\n        def filter_pairs(pairs, drop_gu=False, drop_frac=0.20, seed=0):\n            if not drop_gu:\n                return pairs\n            rng = np.random.RandomState(seed)\n            gu = [(i,j) for (i,j) in pairs if is_gu(seq[i], seq[j])]\n            keep = set(pairs)\n            if gu:\n                rng.shuffle(gu)\n                kdrop = int(len(gu)*drop_frac)\n                for p in gu[:kdrop]:\n                    if p in keep: keep.remove(p)\n            return sorted(list(keep))\n\n        pair_sets = []\n        pair_sets.append(pairs0)\n        pair_sets.append(nussinov_stacking(seq, min_loop=2))\n        pair_sets.append(nussinov_stacking(seq, min_loop=4))\n        pair_sets.append(filter_pairs(pairs0, drop_gu=True, drop_frac=0.35, seed=base_seed+11+r))\n        # extra: perturb by limiting span (bias toward local helices)\n        pairs_span = [(i,j) for (i,j) in pairs0 if (j-i) <= bud[\"max_span\"]]\n        pair_sets.append(pairs_span if pairs_span else pairs0)\n\n        structs = []\n        for k in range(n_structures):\n            seed = base_seed + 100000*r + 97*k\n            coords_k, Ek = mc_two_stage(seq, pair_sets[k], seed, bud)\n            structs.append(coords_k)\n\n        for i, nt in enumerate(seq, start=1):\n            row = {\"ID\": f\"{sid}_{i}\", \"resname\": nt, \"resid\": i}\n            for k in range(n_structures):\n                row[f\"x_{k+1}\"] = float(structs[k][i-1, 0])\n                row[f\"y_{k+1}\"] = float(structs[k][i-1, 1])\n                row[f\"z_{k+1}\"] = float(structs[k][i-1, 2])\n            rows.append(row)\n\n        print(f\"[{r+1}/{len(test_df)}] {sid}: N={N} pairs0={len(pairs0)} steps={bud['steps1']+bud['steps2']} nested={bud.get('use_nested', False)}\")\n\n    return pd.DataFrame(rows)\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T09:07:25.245701Z","iopub.execute_input":"2026-01-22T09:07:25.246029Z","iopub.status.idle":"2026-01-22T09:07:25.256007Z","shell.execute_reply.started":"2026-01-22T09:07:25.246001Z","shell.execute_reply":"2026-01-22T09:07:25.255192Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# -----------------------------\n# 7) Run\n# -----------------------------\nif __name__ == \"__main__\":\n    test_path = \"/kaggle/input/stanford-rna-3d-folding-2/test_sequences.csv\"\n    assert os.path.exists(test_path), f\"{test_path} not found.\"\n    test_df = read_test_sequences(test_path)\n\n    sub = make_submission(test_df, n_structures=5, base_seed=2026)\n    sub.to_csv(\"submission.csv\", index=False)\n    print(\"Wrote submission.csv:\", sub.shape)\n    print(sub.head(3))","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-22T09:07:29.556704Z","iopub.execute_input":"2026-01-22T09:07:29.557296Z","execution_failed":"2026-01-22T09:08:38.095Z"}},"outputs":[],"execution_count":null}]}