{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.11","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":12024591,"sourceType":"competition"},{"sourceId":11676201,"sourceType":"datasetVersion","datasetId":7328231}],"dockerImageVersionId":31011,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        os.path.join(dirname, filename)\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-05-05T18:49:23.261112Z","iopub.execute_input":"2025-05-05T18:49:23.261408Z","iopub.status.idle":"2025-05-05T18:49:25.455975Z","shell.execute_reply.started":"2025-05-05T18:49:23.261383Z","shell.execute_reply":"2025-05-05T18:49:25.455206Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install viennarna\n!pip install forgi\n!pip install LinearFold\n!pip install --no-deps pydca\n!pip install biopython\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T18:49:25.817967Z","iopub.execute_input":"2025-05-05T18:49:25.818719Z","iopub.status.idle":"2025-05-05T18:50:00.601279Z","shell.execute_reply.started":"2025-05-05T18:49:25.818668Z","shell.execute_reply":"2025-05-05T18:50:00.600559Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nlab_df1 = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_labels.csv\")\n#lab_df2 = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_labels.v2.csv\")\n#lab_df = pd.concat([lab_df1, lab_df2], ignore_index=True)\nlab_df = lab_df1\nprint(lab_df.shape)\nprint(lab_df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T18:50:00.602815Z","iopub.execute_input":"2025-05-05T18:50:00.603048Z","iopub.status.idle":"2025-05-05T18:50:00.876978Z","shell.execute_reply.started":"2025-05-05T18:50:00.603028Z","shell.execute_reply":"2025-05-05T18:50:00.876236Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nseq_df1 = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_sequences.csv\")\n#seq_df2 = pd.read_csv(\"/kaggle/input/stanford-rna-3d-folding/train_sequences.v2.csv\")\n#seq_df = pd.concat([seq_df1, seq_df2], ignore_index=True)\nseq_df = seq_df1\nprint(seq_df.shape)\nprint(seq_df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T18:50:04.502644Z","iopub.execute_input":"2025-05-05T18:50:04.503393Z","iopub.status.idle":"2025-05-05T18:50:04.566509Z","shell.execute_reply.started":"2025-05-05T18:50:04.503367Z","shell.execute_reply":"2025-05-05T18:50:04.565869Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"seq_df = pd.read_csv('/kaggle/input/seq-df/seq_df.csv')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T18:50:07.858437Z","iopub.execute_input":"2025-05-05T18:50:07.859132Z","iopub.status.idle":"2025-05-05T18:50:08.086303Z","shell.execute_reply.started":"2025-05-05T18:50:07.859103Z","shell.execute_reply":"2025-05-05T18:50:08.085527Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# seq_df = seq_df.iloc[:200].reset_index(drop=True)\n\nprint(seq_df.shape)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T18:50:09.825317Z","iopub.execute_input":"2025-05-05T18:50:09.825829Z","iopub.status.idle":"2025-05-05T18:50:09.831071Z","shell.execute_reply.started":"2025-05-05T18:50:09.825803Z","shell.execute_reply":"2025-05-05T18:50:09.830219Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Mutual Information","metadata":{}},{"cell_type":"code","source":"import os\nimport math\nimport numpy as np\nimport pandas as pd\nfrom Bio import AlignIO\nfrom multiprocessing import Pool, cpu_count\nfrom tqdm import tqdm\nfrom numba import njit, prange\n\n@njit(parallel=True)\ndef _mi_numba(A, q):\n    N, L = A.shape\n    p = np.zeros((L, q))\n    for i in prange(L):\n        for n in range(N):\n            p[i, A[n, i]] += 1\n        for sym in range(q):\n            p[i, sym] /= N\n\n    mi = np.zeros((L, L))\n    for i in prange(L):\n        for j in range(i, L):\n            joint = np.zeros((q, q))\n            for n in range(N):\n                joint[A[n, i], A[n, j]] += 1\n            for a in range(q):\n                for b in range(q):\n                    joint[a, b] /= N\n\n            m_ij = 0.0\n            for a in range(q):\n                for b in range(q):\n                    pij = joint[a, b]\n                    if pij > 0.0:\n                        m_ij += pij * math.log2(pij / (p[i, a] * p[j, b]))\n            mi[i, j] = m_ij\n            mi[j, i] = m_ij\n    return mi\n\ndef compute_mi_numba(tid):\n    \"\"\"Load MSA for tid, encode it, and return (tid, mi_matrix)\"\"\"\n    path = f\"/kaggle/input/stanford-rna-3d-folding/MSA/{tid}.MSA.fasta\"\n    aln  = AlignIO.read(path, \"fasta\")\n    seqs = [str(rec.seq) for rec in aln]\n    alphabet = sorted({ch for s in seqs for ch in s})\n    sym2i    = {s: i for i, s in enumerate(alphabet)}\n    q        = len(alphabet)\n    A = np.array([[sym2i[ch] for ch in s] for s in seqs], dtype=np.int64)\n\n    mi = _mi_numba(A, q)\n    return tid, mi\n\nnproc = max(1, cpu_count())\nwith Pool(nproc) as pool:\n    results = list(tqdm(\n        pool.imap_unordered(compute_mi_numba, seq_df.target_id),\n        total=len(seq_df),\n        desc=\"Computing MI\"\n    ))\n\nmi_dict = dict(results)\nseq_df['mi_matrix'] = seq_df.target_id.map(mi_dict)\n\nprint(seq_df.loc[:, ['target_id', 'mi_matrix']].head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T14:30:40.784059Z","iopub.execute_input":"2025-05-05T14:30:40.784389Z","iopub.status.idle":"2025-05-05T14:31:30.643339Z","shell.execute_reply.started":"2025-05-05T14:30:40.784342Z","shell.execute_reply":"2025-05-05T14:31:30.642104Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Direct Coupling Analysis","metadata":{}},{"cell_type":"code","source":"# install JAX + CUDA (on Kaggle you usually get it by default; otherwise pip install --upgrade \"jax[cuda]==0.4.14\" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html\nimport os\nimport numpy as np\nimport pandas as pd\nfrom Bio import AlignIO\nimport jax\nimport jax.numpy as jnp\nfrom jax import jit\nfrom tqdm.auto import tqdm\n\ndef load_msa(target_id):\n    path = f\"/kaggle/input/stanford-rna-3d-folding/MSA/{target_id}.MSA.fasta\"\n    aln  = AlignIO.read(path, \"fasta\")\n    seqs = [str(rec.seq) for rec in aln]\n    alpha = {'A':0,'C':1,'G':2,'U':3,'-':4}\n    M = np.array([[alpha.get(r,4) for r in seq] for seq in seqs], dtype=np.int32)\n    return M\n\n@jit\ndef mf_dca_jax(seq_matrix, theta=0.8, H=0.5, lambda_reg=0.01):\n    \"\"\"\n    seq_matrix: int32 array (N, L) with values in [0..4]\n    returns couplings (L, L) on GPU\n    \"\"\"\n    N, L = seq_matrix.shape\n    q = 5\n    # one-hot: (N, L, q)\n    oh = jax.nn.one_hot(seq_matrix, q)\n\n    eq = jnp.einsum(\"n l a, m l a->n m\", oh, oh) / L\n    w = 1.0 / jnp.sum(eq >= theta, axis=1)\n    M_eff = jnp.sum(w)\n\n    f_i = jnp.einsum(\"n, n l a->l a\", w, oh)\n    f_i = (f_i + H/q) / (M_eff + H)\n\n    f_ij = jnp.einsum(\"n, n i a, n j b->i j a b\", w, oh, oh)\n    f_ij = (f_ij + H/(q*q)) / (M_eff + H)\n\n    diff = f_ij - jnp.einsum(\"i a, j b->i j a b\", f_i, f_i)\n    C = diff.reshape((L*q, L*q))\n    C = C + lambda_reg * jnp.eye(L*q)\n\n    invC = jnp.linalg.inv(C)\n    \n    invC4 = invC.reshape((L, q, L, q))\n    # square & sum over a,b; then sqrt\n    J = jnp.sqrt(jnp.sum(invC4**2, axis=(1,3)))\n    return J\n\n_ = mf_dca_jax(jnp.zeros((2,2), dtype=jnp.int32))\n\nresults = []\nfor tid in tqdm(seq_df.target_id, desc=\"MF-DCA (GPU)\"):\n    M = load_msa(tid)\n    J = mf_dca_jax(M)\n    results.append((tid, np.array(J)))  \n\ndca_dict = dict(results)\nseq_df[\"dca\"] = seq_df.target_id.map(dca_dict)\n\ntid0 = seq_df.loc[0, \"target_id\"]\nprint(tid0, seq_df.loc[0, \"dca\"].shape)\nprint(seq_df.loc[:, ['target_id', 'dca']].head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T18:50:20.033783Z","iopub.execute_input":"2025-05-05T18:50:20.034359Z","iopub.status.idle":"2025-05-05T18:51:27.253263Z","shell.execute_reply.started":"2025-05-05T18:50:20.034334Z","shell.execute_reply":"2025-05-05T18:51:27.252664Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Covariation Score","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom Bio import AlignIO\nfrom tqdm.auto import tqdm\nimport multiprocessing as mp\n\n_base2num = {'A':0, 'C':1, 'G':2, 'U':3}\n\ndef compute_covariance(target_id):\n    \"\"\"\n    For a given target_id, load its MSA and compute\n    the L×L covariance matrix on numeric‐encoded columns\n    \"\"\"\n    msa_path = f\"/kaggle/input/stanford-rna-3d-folding/MSA/{target_id}.MSA.fasta\"\n    aln = AlignIO.read(msa_path, \"fasta\")\n\n    seqs = [str(rec.seq) for rec in aln]\n    if len(seqs) == 0:\n        return target_id, None\n\n    numeric_rows = []\n    L = len(seqs[0])\n    for s in seqs:\n        try:\n            row = [_base2num[b] for b in s] \n        except KeyError:\n            continue\n        numeric_rows.append(row)\n\n    M = np.array(numeric_rows, dtype=float)\n    if M.shape[0] < 2:\n        return target_id, np.zeros((L, L), dtype=float)\n\n    Mc = M - M.mean(axis=0, keepdims=True)\n    cov = (Mc.T @ Mc) / (M.shape[0] - 1)\n    return target_id, cov\n\nids = seq_df[\"target_id\"].tolist()\nwith mp.Pool(mp.cpu_count()) as pool:\n    results = list(\n        tqdm(pool.imap_unordered(compute_covariance, ids),\n             total=len(ids),\n             desc=\"Computing covariance\")\n    )\n\ncov_dict = {tid: cov for tid, cov in results}\nseq_df[\"covariance\"] = seq_df[\"target_id\"].map(cov_dict)\n\nprint(seq_df[[\"target_id\",\"covariance\"]].head())\nprint(\"Example shape:\", seq_df.loc[0,\"covariance\"].shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T14:39:45.892913Z","iopub.execute_input":"2025-05-05T14:39:45.893299Z","iopub.status.idle":"2025-05-05T14:39:47.850178Z","shell.execute_reply.started":"2025-05-05T14:39:45.893270Z","shell.execute_reply":"2025-05-05T14:39:47.849330Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Base Pair Distance Potentials","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nfrom multiprocessing import Pool, cpu_count\n\nlab_df[\"target_id\"] = lab_df[\"ID\"].str.rsplit(pat=\"_\", n=1).str[0]\n\nlabels = lab_df.dropna(subset=[\"x_1\",\"y_1\",\"z_1\"])\n\nbases = [\"A\",\"C\",\"G\",\"U\"]\nnt2idx = {n:i for i,n in enumerate(bases)}\n\npair_dists = {tuple(sorted((b1,b2))): [] for b1 in bases for b2 in bases}\n\nfor tid, grp in tqdm(labels.groupby(\"target_id\"), desc=\"Gathering distances\"):\n    grp = grp.sort_values(\"resid\")\n    coords = grp[[\"x_1\",\"y_1\",\"z_1\"]].to_numpy()\n    resns  = grp[\"resname\"].to_numpy()\n    dmat = np.linalg.norm(coords[:,None,:] - coords[None,:,:], axis=2)\n    L = len(resns)\n    for i in range(L):\n        b1 = resns[i]\n        if b1 not in nt2idx: \n            continue\n        for j in range(i+1, L):\n            b2 = resns[j]\n            if b2 not in nt2idx:\n                continue\n            pair = tuple(sorted((b1, b2)))\n            pair_dists[pair].append(dmat[i, j])\n\nbp_mean = np.full((4,4), np.nan, dtype=np.float32)\nfor (b1, b2), dist_list in pair_dists.items():\n    if dist_list:\n        m = float(np.mean(dist_list))\n        i, j = nt2idx[b1], nt2idx[b2]\n        bp_mean[i, j] = bp_mean[j, i] = m\n\nprint(\"bp_mean (rows/cols = A,C,G,U):\\n\", bp_mean)\n\ndef compute_bp_potentials(seq: str):\n    \"\"\"\n    Build an (L,L) matrix where entry (i,j) = bp_mean[base_i, base_j],\n    skipping any non-ACGU bases\n    \"\"\"\n    L = len(seq)\n    mat = np.zeros((L, L), dtype=np.float32)\n    for i, b1 in enumerate(seq):\n        i1 = nt2idx.get(b1)\n        if i1 is None:\n            continue\n        for j, b2 in enumerate(seq):\n            i2 = nt2idx.get(b2)\n            if i2 is None:\n                continue\n            mat[i, j] = bp_mean[i1, i2]\n    return mat\n\nwith Pool(cpu_count()) as pool:\n    bp_mats = list(\n        tqdm(pool.imap(compute_bp_potentials, seq_df.sequence),\n             total=len(seq_df),\n             desc=\"Computing bp_potentials\")\n    )\n\nseq_df[\"bp_potentials\"] = bp_mats\n\nprint(seq_df.loc[:, ['target_id', 'bp_potentials']].head()) \n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T14:57:29.037272Z","iopub.execute_input":"2025-05-05T14:57:29.037634Z","iopub.status.idle":"2025-05-05T14:59:41.308341Z","shell.execute_reply.started":"2025-05-05T14:57:29.037611Z","shell.execute_reply":"2025-05-05T14:59:41.307437Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Euclidean Ground Truth Distances","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport math\nfrom numba import njit, prange\n\nlab_df[\"target_id\"] = lab_df[\"ID\"].str.rsplit(pat=\"_\", n=1).str[0]\n\ncoords_dict = {}\nfor tid, group in lab_df.groupby(\"target_id\"):\n    grp = group.sort_values(\"resid\")\n    coords_dict[tid] = grp[[\"x_1\",\"y_1\",\"z_1\"]].to_numpy(dtype=np.float64)\n\n@njit(parallel=True, cache=True)\ndef compute_distances_jit(coords):\n    L = coords.shape[0]\n    dist = np.empty((L, L), dtype=np.float64)\n    for i in prange(L):\n        xi, yi, zi = coords[i,0], coords[i,1], coords[i,2]\n        for j in range(L):\n            dx = xi - coords[j,0]\n            dy = yi - coords[j,1]\n            dz = zi - coords[j,2]\n            dist[i, j] = math.sqrt(dx*dx + dy*dy + dz*dz)\n    return dist\n\n_ = compute_distances_jit(np.zeros((2,3), dtype=np.float64))\n\ngt_distances = { tid: compute_distances_jit(coords)\n                 for tid, coords in coords_dict.items() }\n\nseq_df[\"gt_distances\"] = seq_df[\"target_id\"].map(gt_distances)\n\nprint(seq_df[[\"target_id\",\"gt_distances\"]].head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T15:07:09.386410Z","iopub.execute_input":"2025-05-05T15:07:09.387582Z","iopub.status.idle":"2025-05-05T15:07:21.143081Z","shell.execute_reply.started":"2025-05-05T15:07:09.387545Z","shell.execute_reply":"2025-05-05T15:07:21.142355Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Binary Contact Map","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom numba import jit\nfrom tqdm import tqdm\nimport multiprocessing as mp\n\nlab_df[\"target_id\"] = lab_df[\"ID\"].str.rsplit(pat=\"_\", n=1).str[0]\ncoords_df = lab_df[['target_id','resid','x_1','y_1','z_1']]\ngrouped = coords_df.groupby('target_id')\n\n@jit(nopython=True)\ndef make_contact_map(coords: np.ndarray, thresh: float) -> np.ndarray:\n    L = coords.shape[0]\n    out = np.zeros((L, L), dtype=np.uint8)\n    thr2 = thresh * thresh\n    for i in range(L):\n        xi, yi, zi = coords[i]\n        for j in range(L):\n            dx = xi - coords[j,0]\n            dy = yi - coords[j,1]\n            dz = zi - coords[j,2]\n            if dx*dx + dy*dy + dz*dz <= thr2:\n                out[i, j] = 1\n    return out\n\ndef process_target(tid, threshold=8.0):\n    grp = grouped.get_group(tid).sort_values('resid')\n    coords = grp[['x_1','y_1','z_1']].values.astype(np.float64)\n    cmap = make_contact_map(coords, threshold)\n    return tid, cmap\n\ntargets = seq_df['target_id'].tolist()\nthreshold = 8.0\n\nwith mp.Pool(processes=mp.cpu_count()) as pool:\n    args = [(t, threshold) for t in targets]\n    results = list(tqdm(pool.starmap(process_target, args),\n                        total=len(args),\n                        desc=\"Computing contact maps\"))\n\ncontact_dict = dict(results)\nseq_df['contact_map'] = seq_df['target_id'].map(contact_dict)\n\nprint(seq_df[['target_id','contact_map']].head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T15:12:28.033606Z","iopub.execute_input":"2025-05-05T15:12:28.033933Z","iopub.status.idle":"2025-05-05T15:12:35.445742Z","shell.execute_reply.started":"2025-05-05T15:12:28.033907Z","shell.execute_reply":"2025-05-05T15:12:35.444478Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Angle/Dihedral Matrix","metadata":{}},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nimport numba\nfrom concurrent.futures import ProcessPoolExecutor, as_completed\nfrom tqdm.auto import tqdm\n\ncoords_dict = {}\nfor tid, group in lab_df.groupby(lab_df[\"ID\"].str.rsplit(pat=\"_\", n=1).str[0]):\n    g = group.sort_values('resid')\n    coords = g[['x_1','y_1','z_1']].to_numpy(dtype=np.float64)\n    coords_dict[tid] = coords\n\n@numba.njit(parallel=True)\ndef compute_angle_matrix(coords):\n    \"\"\"\n    Given coords[L,3], return angles[L,L] where\n    angles[i,j] = angle between backbone vector at i and vector from i to j.\n    \"\"\"\n    L = coords.shape[0]\n    mat = np.zeros((L, L), dtype=np.float64)\n    for i in numba.prange(L):\n        if i < L-1:\n            vi = coords[i+1] - coords[i]\n        else:\n            vi = coords[i] - coords[i-1]\n        n_vi = np.sqrt(vi[0]**2 + vi[1]**2 + vi[2]**2) + 1e-8\n        for j in range(L):\n            r = coords[j] - coords[i]\n            n_r = np.sqrt(r[0]**2 + r[1]**2 + r[2]**2) + 1e-8\n            dot = vi[0]*r[0] + vi[1]*r[1] + vi[2]*r[2]\n            c = dot / (n_vi * n_r)\n            if c > 1: c = 1\n            elif c < -1: c = -1\n            mat[i, j] = np.arccos(c)\n    return mat\n\ndef worker(tid):\n    coords = coords_dict[tid]\n    return tid, compute_angle_matrix(coords)\n\nresults = {}\nwith ProcessPoolExecutor() as exe:\n    futures = [exe.submit(worker, tid) for tid in seq_df.target_id]\n    for fut in tqdm(as_completed(futures), total=len(futures), desc=\"Angle matrices\"):\n        tid, ang = fut.result()\n        results[tid] = ang\n\nseq_df['angle_matrix'] = seq_df['target_id'].map(results)\n\nprint(seq_df[['target_id','sequence','angle_matrix']].head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-05-05T15:17:35.095579Z","iopub.execute_input":"2025-05-05T15:17:35.096169Z","iopub.status.idle":"2025-05-05T15:17:42.850508Z","shell.execute_reply.started":"2025-05-05T15:17:35.096145Z","shell.execute_reply":"2025-05-05T15:17:42.849467Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"seq_df.to_csv(\"seq_df_2D.csv\", index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}