{"cells":[{"cell_type":"code","execution_count":null,"metadata":{},"outputs":[],"source":"# Stanford RNA 3D Folding 2 — Sequence-aware Template-Based Modeling (TBM)\n# CPU-only, deterministic. Adapted from clean public TBM reference\n# (analyticaobscura/rna-3d-folding-tbm) using the competition's own train data.\nimport subprocess, sys, glob, os\n\n# Install biopython from offline wheel dataset (recursive find = mount-path agnostic)\ntry:\n    import Bio  # may already be present in the image\nexcept ImportError:\n    _whls = glob.glob('/kaggle/input/**/biopython*.whl', recursive=True)\n    print('biopython wheels found:', _whls)\n    if not _whls:\n        print('--- .whl files under /kaggle/input ---')\n        for _r, _d, _f in os.walk('/kaggle/input'):\n            for _fn in _f:\n                if _fn.endswith('.whl'):\n                    print(os.path.join(_r, _fn))\n        raise FileNotFoundError('biopython wheel not found under /kaggle/input')\n    subprocess.run([sys.executable, '-m', 'pip', 'install', '--no-index', '--no-deps', _whls[0]], check=True)\n\nimport numpy as np\nimport pandas as pd\nimport random, hashlib, time, warnings\nfrom pathlib import Path\nfrom scipy.spatial import distance_matrix\nfrom scipy.spatial.transform import Rotation as R\nwarnings.filterwarnings('ignore')\nfrom Bio.Align import PairwiseAligner\n\nSEED = 21\nnp.random.seed(SEED); random.seed(SEED)\n\n# Global RNA-tuned aligner\nALIGNER = PairwiseAligner()\nALIGNER.mode = 'global'\nALIGNER.match_score = 2.5\nALIGNER.mismatch_score = -1\nALIGNER.open_gap_score = -10\nALIGNER.extend_gap_score = -0.5\n\n# ---- Data path detection ----\nROOT = Path('/kaggle/input')\nDATA_DIR = None\nfor p in ROOT.rglob('test_sequences.csv'):\n    DATA_DIR = p.parent\n    break\nif DATA_DIR is None:\n    raise FileNotFoundError('test_sequences.csv not found')\nprint('DATA_DIR:', DATA_DIR)\n\ntest_seqs = pd.read_csv(DATA_DIR / 'test_sequences.csv')\ntrain_seqs = pd.read_csv(DATA_DIR / 'train_sequences.csv')\ntrain_labels = pd.read_csv(DATA_DIR / 'train_labels.csv', low_memory=False)\nsample_sub = pd.read_csv(DATA_DIR / 'sample_submission.csv')\nprint('train_seqs', train_seqs.shape, '| test_seqs', test_seqs.shape, '| sample_sub', sample_sub.shape)\n\nSEQ_COL = 'sequence' if 'sequence' in test_seqs.columns else test_seqs.columns[1]\nTID_COL = 'target_id' if 'target_id' in train_seqs.columns else train_seqs.columns[0]\n\n# ---- Build train coord library ----\ndef process_labels(labels_df):\n    df = labels_df.copy()\n    df['target_id'] = df['ID'].str.rsplit('_', n=1).str[0]\n    df = df.sort_values(['target_id', 'resid'])\n    cc = ['x_1', 'y_1', 'z_1']\n    arr = df[cc].values.astype(np.float64).copy()\n    arr[arr < -1e6] = np.nan\n    df[cc] = arr\n    d = {}\n    for tid, g in df.groupby('target_id', sort=False):\n        d[tid] = g[cc].values\n    return d\n\ntrain_coords_dict = process_labels(train_labels)\nprint('train structures:', len(train_coords_dict))\n\n# Keep only train seqs that have coords\ntrain_seqs = train_seqs[train_seqs[TID_COL].isin(train_coords_dict)].reset_index(drop=True)\nprint('usable train templates:', len(train_seqs))\n\n# ---- Alignment helpers ----\ndef get_alignment_score(query_seq, template_seq):\n    al = ALIGNER.align(query_seq, template_seq)\n    best = next(iter(al), None)\n    if best is None:\n        return 0.0\n    max_possible = 2.1 * min(len(query_seq), len(template_seq))\n    return min(best.score / max_possible, 1.0)\n\ndef build_aligned_sequences(query_seq, template_seq, alignment):\n    qr, tr = alignment.aligned\n    if len(qr) == 0:\n        return None, None\n    aq, at = [], []\n    qp = tp = 0\n    for (qs, qe), (ts, te) in zip(qr, tr):\n        while qp < qs:\n            aq.append(query_seq[qp]); at.append('-'); qp += 1\n        while tp < ts:\n            aq.append('-'); at.append(template_seq[tp]); tp += 1\n        for i in range(qe - qs):\n            aq.append(query_seq[qs + i]); at.append(template_seq[ts + i])\n        qp, tp = qe, te\n    while qp < len(query_seq):\n        aq.append(query_seq[qp]); at.append('-'); qp += 1\n    while tp < len(template_seq):\n        aq.append('-'); at.append(template_seq[tp]); tp += 1\n    return ''.join(aq), ''.join(at)\n\ndef get_aligned_sequences(query_seq, template_seq):\n    al = ALIGNER.align(query_seq, template_seq)\n    best = next(iter(al), None)\n    if best is None:\n        return None, None\n    return build_aligned_sequences(query_seq, template_seq, best)\n\ndef find_similar_sequences(query_seq, top_n=5):\n    out = []\n    qlen = len(query_seq)\n    for tid, tseq in TRAIN_SEQ_PAIRS:\n        d = abs(len(tseq) - qlen) / max(len(tseq), qlen)\n        if d > 0.5:\n            continue\n        s = get_alignment_score(query_seq, tseq)\n        out.append((tid, tseq, s))\n    out.sort(key=lambda x: x[2], reverse=True)\n    return out[:top_n]\n\ndef fill_coordinate_gaps(coords):\n    n = len(coords); step = 4.0\n    for i in range(n):\n        if np.isnan(coords[i, 0]):\n            pv = next((j for j in range(i - 1, -1, -1) if not np.isnan(coords[j, 0])), -1)\n            nv = next((j for j in range(i + 1, n) if not np.isnan(coords[j, 0])), -1)\n            if pv >= 0 and nv >= 0:\n                w = (i - pv) / (nv - pv)\n                coords[i] = (1 - w) * coords[pv] + w * coords[nv]\n    for i in range(n):\n        if np.isnan(coords[i, 0]):\n            if i == 0:\n                fv = next((j for j in range(1, n) if not np.isnan(coords[j, 0])), -1)\n                if fv >= 0:\n                    for j in range(fv - 1, -1, -1):\n                        dd = np.random.normal(0, 1, 3)\n                        dd = dd / (np.linalg.norm(dd) + 1e-10) * step\n                        coords[j] = coords[j + 1] - dd\n                else:\n                    coords = generate_basic_structure_coords(n); break\n            else:\n                pv = next((j for j in range(i - 1, -1, -1) if not np.isnan(coords[j, 0])), -1)\n                if pv >= 0:\n                    dd = np.random.normal(0, 1, 3)\n                    dd = dd / (np.linalg.norm(dd) + 1e-10) * step\n                    coords[i] = coords[pv] + dd\n    return np.nan_to_num(coords)\n\ndef generate_basic_structure_coords(n):\n    c = np.zeros((n, 3))\n    for i in range(n):\n        a = i * 0.6\n        c[i] = [10.0 * np.cos(a), 10.0 * np.sin(a), i * 2.5]\n    return c\n\ndef adapt_template_simple(query_seq, template_seq, tcoords):\n    qc = np.zeros((len(query_seq), 3))\n    scale = len(tcoords) / len(query_seq)\n    for i in range(len(query_seq)):\n        ti = min(int(i * scale), len(tcoords) - 1)\n        qc[i] = tcoords[ti] if not np.any(np.isnan(tcoords[ti])) else [np.nan, np.nan, np.nan]\n    return fill_coordinate_gaps(qc)\n\ndef adapt_template_to_query(query_seq, template_seq, tcoords):\n    aq, at = get_aligned_sequences(query_seq, template_seq)\n    if aq is None:\n        return adapt_template_simple(query_seq, template_seq, tcoords)\n    qc = np.full((len(query_seq), 3), np.nan)\n    qi = ti = 0\n    for i in range(len(aq)):\n        qch, tch = aq[i], at[i]\n        if qch != '-' and tch != '-':\n            if qi < len(query_seq) and ti < len(tcoords) and not np.any(np.isnan(tcoords[ti])):\n                qc[qi] = tcoords[ti]\n            ti += 1; qi += 1\n        elif qch != '-' and tch == '-':\n            qi += 1\n        elif qch == '-' and tch != '-':\n            ti += 1\n    return fill_coordinate_gaps(qc)\n\ndef adaptive_rna_constraints(coords, sequence, confidence=1.0):\n    rc = coords.copy(); n = len(sequence)\n    cs = 0.8 * (1.0 - min(confidence, 0.8))\n    lo, hi = 5.5, 6.5\n    for i in range(n - 1):\n        cur, nxt = rc[i], rc[i + 1]\n        dist = np.linalg.norm(nxt - cur)\n        if dist < lo or dist > hi:\n            tgt = (lo + hi) / 2\n            dr = nxt - cur\n            dr = dr / (np.linalg.norm(dr) + 1e-10)\n            adj = (tgt - dist) * cs\n            rc[i + 1] = cur + dr * (dist + adj)\n    mind = 3.8\n    dm = distance_matrix(rc, rc)\n    cl = np.where((dm < mind) & (dm > 0))\n    for k in range(len(cl[0])):\n        i, j = cl[0][k], cl[1][k]\n        if abs(i - j) <= 1 or i >= j:\n            continue\n        pi, pj = rc[i], rc[j]\n        dr = pj - pi\n        dr = dr / (np.linalg.norm(dr) + 1e-10)\n        adj = (mind - dm[i, j]) * cs\n        rc[i] = pi - dr * (adj / 2)\n        rc[j] = pj + dr * (adj / 2)\n    return rc\n\ndef generate_rna_structure(sequence, seed=None):\n    if seed is not None:\n        np.random.seed(seed); random.seed(seed)\n    n = len(sequence)\n    c = np.zeros((n, 3))\n    for i in range(min(3, n)):\n        a = i * 0.6\n        c[i] = [10.0 * np.cos(a), 10.0 * np.sin(a), i * 2.5]\n    cd = np.array([0.0, 0.0, 1.0])\n    comp = {'G': 'C', 'C': 'G', 'A': 'U', 'U': 'A'}\n    for i in range(3, n):\n        base = sequence[i]; paired = False; pidx = -1\n        w = min(i, 15)\n        for j in range(i - w, i):\n            if j >= 0 and sequence[j] == comp.get(base, 'X'):\n                paired = True; pidx = j; break\n        if paired and i - pidx <= 10 and random.random() < 0.7:\n            pp = c[pidx]\n            off = np.random.normal(0, 1, 3) * 2.0\n            bpd = 10.0 + random.uniform(-1.0, 1.0)\n            center = np.mean(c[:i], axis=0)\n            dr = center - pp\n            dr = dr / (np.linalg.norm(dr) + 1e-10)\n            c[i] = pp + dr * bpd + off\n            cd = np.random.normal(0, 0.3, 3)\n            cd = cd / (np.linalg.norm(cd) + 1e-10)\n        else:\n            if random.random() < 0.3:\n                ang = random.uniform(0.2, 0.6)\n                ax = np.random.normal(0, 1, 3)\n                ax = ax / (np.linalg.norm(ax) + 1e-10)\n                cd = R.from_rotvec(ang * ax).apply(cd)\n            else:\n                cd += np.random.normal(0, 0.15, 3)\n                cd = cd / (np.linalg.norm(cd) + 1e-10)\n            c[i] = c[i - 1] + random.uniform(3.5, 4.5) * cd\n    return c\n\ndef predict_rna_structures(sequence, target_id, n_predictions=5):\n    preds = []\n    sims = find_similar_sequences(sequence, top_n=n_predictions)\n    for tid, tseq, sim in sims:\n        ac = adapt_template_to_query(sequence, tseq, train_coords_dict[tid])\n        rc = adaptive_rna_constraints(ac, sequence, confidence=sim)\n        scale = max(0.05, 0.8 - sim)\n        preds.append(rc + np.random.normal(0, scale, rc.shape))\n        if len(preds) >= n_predictions:\n            break\n    while len(preds) < n_predictions:\n        us = f\"{target_id}_{len(preds)}\"\n        sv = int(hashlib.sha256(us.encode()).hexdigest(), 16) % (2**32 - 1)\n        dn = generate_rna_structure(sequence, seed=sv)\n        preds.append(adaptive_rna_constraints(dn, sequence, confidence=0.2))\n    return preds[:n_predictions]\n\n# ---- Build predictions aligned to sample_submission ----\nTRAIN_SEQ_PAIRS = list(zip(train_seqs[TID_COL].astype(str), train_seqs['sequence'].astype(str)))\ntest_seq_map = dict(zip(test_seqs[TID_COL if TID_COL in test_seqs.columns else 'target_id'].astype(str),\n                        test_seqs[SEQ_COL].astype(str)))\n\nsubmission = sample_sub.copy()\nsubmission['_target'] = submission['ID'].str.rsplit('_', n=1).str[0]\n\nt0 = time.time()\ngroups = list(submission.groupby('_target', sort=False))\nprint('test targets:', len(groups))\nfor gi, (target_id, group) in enumerate(groups):\n    idx = group.index\n    seq = test_seq_map.get(str(target_id))\n    if seq is None or len(seq) != len(group):\n        seq = 'A' * len(group)\n    rs = int(hashlib.sha256(str(target_id).encode()).hexdigest(), 16) % (2**32 - 1)\n    np.random.seed(rs); random.seed(rs)\n    preds = predict_rna_structures(seq, str(target_id), n_predictions=5)\n    for s in range(5):\n        submission.loc[idx, f'x_{s+1}'] = np.round(preds[s][:, 0], 3)\n        submission.loc[idx, f'y_{s+1}'] = np.round(preds[s][:, 1], 3)\n        submission.loc[idx, f'z_{s+1}'] = np.round(preds[s][:, 2], 3)\n    if (gi + 1) % 10 == 0:\n        print(f'  {gi+1}/{len(groups)} targets done, {time.time()-t0:.0f}s')\n\nsubmission = submission.drop(columns=['_target'])\nout_path = '/kaggle/working/submission.csv'\nsubmission.to_csv(out_path, index=False)\nprint('Saved', out_path, submission.shape)\nprint(submission.head())\n"}],"metadata":{"kernelspec":{"name":"python3","display_name":"Python 3","language":"python"},"language_info":{"name":"python"}},"nbformat":4,"nbformat_minor":5}