{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"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":"nvidiaTeslaT4","dataSources":[{"sourceId":51294,"databundleVersionId":7331882,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":6233380,"sourceType":"datasetVersion","datasetId":3580819},{"sourceId":6266392,"sourceType":"datasetVersion","datasetId":3580757},{"sourceId":14485988,"sourceType":"datasetVersion","datasetId":9252413},{"sourceId":14486627,"sourceType":"datasetVersion","datasetId":9252784}],"dockerImageVersionId":30528,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# Grab the code and weights -- Shujun saved these in Kaggle 'datasets'\n!rsync -avzL /kaggle/input/rhofold/RhoFold . >/dev/null 2>&1\n!rsync -avzL /kaggle/input/rhofold-pretrained-weights/rhofold_pretrained.pt RhoFold/pretrained/rhofold_pretrained.pt\n# Grab the pre-trained model as the documentation says\n#!wget https://proj.cse.cuhk.edu.hk/aihlab/rhofold/api/download?filename=rhofold_pretrained.pt -O RhoFold/pretrained/rhofold_pretrained.pt\n","metadata":{"execution":{"iopub.status.busy":"2026-01-13T13:33:24.368287Z","iopub.execute_input":"2026-01-13T13:33:24.368706Z","iopub.status.idle":"2026-01-13T13:33:37.152866Z","shell.execute_reply.started":"2026-01-13T13:33:24.368675Z","shell.execute_reply":"2026-01-13T13:33:37.151771Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Setup RhoFold through conda. \nNote that the weights and original code are no longer publicly available, but these files are stored in Kaggle data sets *rhofold* and *rhofold-pretrained-weights*.","metadata":{}},{"cell_type":"code","source":"!conda env create -f /kaggle/working/RhoFold/envs/environment_linux.yaml >/dev/null 2>&1\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T16:30:16.456262Z","iopub.execute_input":"2026-01-13T16:30:16.457354Z","iopub.status.idle":"2026-01-13T16:30:18.483568Z","shell.execute_reply.started":"2026-01-13T16:30:16.457312Z","shell.execute_reply":"2026-01-13T16:30:18.482448Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!source /opt/conda/bin/activate rhofold && python RhoFold/setup.py install","metadata":{"execution":{"iopub.status.busy":"2026-01-13T16:30:42.636028Z","iopub.execute_input":"2026-01-13T16:30:42.636737Z","iopub.status.idle":"2026-01-13T16:30:49.241294Z","shell.execute_reply.started":"2026-01-13T16:30:42.636706Z","shell.execute_reply":"2026-01-13T16:30:49.239880Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Helper functions to wrap RhoFold","metadata":{}},{"cell_type":"code","source":"def get_rhofold_3D(sequence):\n    #sequence='AUGCAUGCAUGC'\n    with open('fasta.fasta','w+') as f:\n        f.write('>test\\n')\n        f.write(sequence)\n\n    ! source /opt/conda/bin/activate rhofold && python RhoFold/inference.py --input_fas fasta.fasta --single_seq_pred True --output_dir ./rhofold_output/ --ckpt ./RhoFold/pretrained/rhofold_pretrained.pt\n    \ndef get_rhofold_3D_GPU(sequence):\n    #sequence='AUGCAUGCAUGC'\n    with open('fasta.fasta','w+') as f:\n        f.write('>test\\n')\n        f.write(sequence)\n\n    ! source /opt/conda/bin/activate rhofold && python RhoFold/inference.py --input_fas fasta.fasta --single_seq_pred True --output_dir ./rhofold_output/ --ckpt ./RhoFold/pretrained/rhofold_pretrained.pt --device cuda:0 #>/dev/null 2>&1","metadata":{"execution":{"iopub.status.busy":"2026-01-13T13:50:47.042946Z","iopub.execute_input":"2026-01-13T13:50:47.043680Z","iopub.status.idle":"2026-01-13T13:50:47.050862Z","shell.execute_reply.started":"2026-01-13T13:50:47.043646Z","shell.execute_reply":"2026-01-13T13:50:47.050094Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Get py3Dmol installed, load up a PDB.","metadata":{}},{"cell_type":"code","source":"!pip install py3Dmol\n!wget --no-check-certificate 'https://docs.google.com/uc?export=download&id=1us10sgCvCAgqL2XVW_Df2gXX8US_qlyz' -O unrelaxed_model_PymolOUT.pdb\nimport py3Dmol","metadata":{"execution":{"iopub.status.busy":"2026-01-13T13:50:58.628973Z","iopub.execute_input":"2026-01-13T13:50:58.629282Z","iopub.status.idle":"2026-01-13T13:51:09.321683Z","shell.execute_reply.started":"2026-01-13T13:50:58.629256Z","shell.execute_reply":"2026-01-13T13:51:09.320820Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Load in Rhofold model. \n","metadata":{}},{"cell_type":"code","source":"def reorder( resi_lines):\n    '''Do backbone atom reordering for a block of lines corresponding to one residue'''\n    #if len(resi_lines) > 0: print( resi_lines[0][19])  # ACGU\n    #backbone_atoms_for_nt = {'A':['P',\"C5'\",\"O5'\",\"C4'\",\"O4'\",\"C3'\",\"O3'\",\"C2'\",\"O2'\",\"C1'\"]}\n    backbone_atoms = ['P',\"C5'\",\"O5'\",\"C4'\",\"O4'\",\"C3'\",\"O3'\",\"C2'\",\"O2'\",\"C1'\",'OP1','OP2']\n    other_lines = []\n    resi_line_dict = {}\n    for resi_line in resi_lines:\n        if len(resi_line)>25: \n            atom_name = resi_line[11:16].strip()\n            if atom_name in backbone_atoms:\n                resi_line_dict[atom_name] = resi_line\n            else:\n                other_lines.append(resi_line)\n        else: other_lines.append(resi_line)\n    resi_lines_reorder = [resi_line_dict[x] for x in backbone_atoms if x in resi_line_dict]\n    resi_lines_reorder += other_lines\n    return resi_lines_reorder\n    \ndef reorder_atom_numbers( lines ):\n    ''' Atom numbers should go from 1,2,.... '''\n    atom_number = 0\n    lines_new = []\n    for line in lines:\n        line_new = line\n        if len(line) > 11:\n            atom_number += 1\n            line_new = '%s%6d%s' % (line[:5],atom_number,line[11:])\n            assert( len(line_new) == len(line))\n        lines_new.append( line_new )\n    return lines_new\n\ndef sanitize( lines ):\n    ''' sanitize RNA residue lines from RhoFold to be in \"correct\" order expected by py3Dmol.'''\n    lines_is_string = False\n    if isinstance(lines,str): \n        lines = lines.split('\\n')\n        lines_is_string = True\n        \n    lines_reorder = []\n    resi_lines = []\n    chain = ''\n    resi = ''\n    count = 0\n    for (i,line) in enumerate(lines):\n        if len(line)>25 and (line[21] != chain or resi != line[22:26]):\n            #print(len(resi_lines),chain,resi)\n            chain = line[21]\n            resi = line[22:26]\n            resi_lines_reorder = reorder(resi_lines)\n            lines_reorder = lines_reorder + reorder(resi_lines)\n            # get next block ready\n            resi_lines = [line]\n            count += 1\n        else:\n            resi_lines.append(line)\n    lines_reorder = lines_reorder + reorder(resi_lines)\n    lines_reorder = reorder_atom_numbers( lines_reorder )\n    if lines_is_string: lines_reorder = '\\n'.join(lines_reorder)\n    return lines_reorder\n\ndef get_pdb_lines(filename):\n    ''' Currently only reads in ATOM, not HETATM '''\n    lines = ''.join([(line) for line in open(pdb_file_path).readlines() if line.find(\"ATOM\")==0])\n    return sanitize(lines)\n    ","metadata":{"execution":{"iopub.status.busy":"2026-01-13T13:51:25.553491Z","iopub.execute_input":"2026-01-13T13:51:25.554409Z","iopub.status.idle":"2026-01-13T13:51:25.565386Z","shell.execute_reply.started":"2026-01-13T13:51:25.554375Z","shell.execute_reply":"2026-01-13T13:51:25.564371Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# A test sequence ('AB50') from Eterna.\nThis sequence, designed by Eterna player AmyBarish, appears to fold into a pseudoknot, based on experimental chemical mapping data. It also is folded into a pseudoknot by RhoFold.","metadata":{}},{"cell_type":"code","source":"# Sequence for AB50\nget_rhofold_3D_GPU('AGCUCGAAAAGGACCAAAACGAGCAAUAACAUAAAUAAAGGUCCAAAUAA')\npdb_file_path = '/kaggle/working/rhofold_output/unrelaxed_model.pdb';\nview = py3Dmol.view(width=400, height=300)\nview.addModel(get_pdb_lines(pdb_file_path),'pdb')\nview.setStyle( {\"cartoon\": {'color': 'spectrum'}})\nview.zoomTo()\nview.show()","metadata":{"execution":{"iopub.status.busy":"2026-01-13T14:06:53.632496Z","iopub.execute_input":"2026-01-13T14:06:53.632868Z","iopub.status.idle":"2026-01-13T14:07:03.044985Z","shell.execute_reply.started":"2026-01-13T14:06:53.632836Z","shell.execute_reply":"2026-01-13T14:07:03.044008Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"# Running on the Das Dataset of  RNA Benchmark!\n## Link of the paper : https://rdcu.be/eYSGa","metadata":{}},{"cell_type":"code","source":"import os, glob, pathlib, subprocess\n\nINPUT_DIR = \"/kaggle/input/rna-fasta-14das-desirna/RhoFold_Input_clean\"\nOUT_ROOT  = \"/kaggle/working/rhofold_Das_14_outputs\"\nCKPT      = \"/kaggle/working/RhoFold/pretrained/rhofold_pretrained.pt\"\nos.makedirs(OUT_ROOT, exist_ok=True)\n","metadata":{"execution":{"iopub.status.busy":"2026-01-13T14:15:08.632977Z","iopub.execute_input":"2026-01-13T14:15:08.633774Z","iopub.status.idle":"2026-01-13T14:15:08.637924Z","shell.execute_reply.started":"2026-01-13T14:15:08.633741Z","shell.execute_reply":"2026-01-13T14:15:08.637090Z"},"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# collect fasta files\nfasta_files = []\nfor ext in (\"*.fasta\", \"*.fa\", \"*.fna\"):\n    fasta_files += glob.glob(os.path.join(INPUT_DIR, ext))\nfasta_files = sorted(fasta_files)\n\nprint(f\"Found {len(fasta_files)} FASTA files\")\nprint(\"Example:\", fasta_files[:3])\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T14:15:16.441695Z","iopub.execute_input":"2026-01-13T14:15:16.442473Z","iopub.status.idle":"2026-01-13T14:15:16.456494Z","shell.execute_reply.started":"2026-01-13T14:15:16.442444Z","shell.execute_reply":"2026-01-13T14:15:16.455751Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"for fasta_path in fasta_files:\n    name = pathlib.Path(fasta_path).stem\n    out_dir = os.path.join(OUT_ROOT, name)\n    os.makedirs(out_dir, exist_ok=True)\n\n    cmd = f\"\"\"\n    source /opt/conda/bin/activate rhofold && \\\n    python /kaggle/working/RhoFold/inference.py \\\n      --input_fas \"{fasta_path}\" \\\n      --single_seq_pred True \\\n      --output_dir \"{out_dir}\" \\\n      --ckpt \"{CKPT}\" \\\n      --device cuda:0\n    \"\"\"\n    print(f\"\\n=== Running: {name} ===\")\n    r = subprocess.run(cmd, shell=True, executable=\"/bin/bash\")\n    if r.returncode != 0:\n        print(f\"[!] Failed on {fasta_path} (return code {r.returncode})\")\n        # continue to next file instead of stopping\n        continue\n\nprint(\"\\nDone. Outputs in:\", OUT_ROOT)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T14:15:35.205893Z","iopub.execute_input":"2026-01-13T14:15:35.206500Z","iopub.status.idle":"2026-01-13T14:16:37.003298Z","shell.execute_reply.started":"2026-01-13T14:15:35.206472Z","shell.execute_reply":"2026-01-13T14:16:37.002365Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  RMSD using kabsch FIT","metadata":{}},{"cell_type":"code","source":"import os, glob, pathlib\nimport numpy as np\nimport pandas as pd\nfrom Bio.PDB import PDBParser, PDBIO\n\nREF_DIR  = \"/kaggle/input/das-14-rnas-benchmarks/Das-14-RNAs\"\nPRED_ROOT = \"/kaggle/working/rhofold_Das_14_outputs\"\nALIGNED_OUT = \"/kaggle/working/rhofold_Das_14_aligned\"\nCSV_OUT = \"/kaggle/working/rhofold_vs_ref_rmsd.csv\"\n\nos.makedirs(ALIGNED_OUT, exist_ok=True)\n\nparser = PDBParser(QUIET=True)\n\nBACKBONE = [\"P\", \"O5'\", \"C5'\", \"C4'\", \"C3'\", \"O3'\"]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T14:48:37.961028Z","iopub.execute_input":"2026-01-13T14:48:37.961401Z","iopub.status.idle":"2026-01-13T14:48:38.059758Z","shell.execute_reply.started":"2026-01-13T14:48:37.961374Z","shell.execute_reply":"2026-01-13T14:48:38.058828Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"### Needed Functions","metadata":{}},{"cell_type":"code","source":"def find_reference_pdb(rna_name: str):\n    \"\"\"Find a reference PDB in REF_DIR matching rna_name (stem match, case-insensitive).\"\"\"\n    # try exact common patterns first\n    for ext in (\"pdb\", \"PDB\"):\n        cand = os.path.join(REF_DIR, f\"{rna_name}.{ext}\")\n        if os.path.exists(cand):\n            return cand\n\n    # else search by stem\n    ref_files = glob.glob(os.path.join(REF_DIR, \"*.pdb\")) + glob.glob(os.path.join(REF_DIR, \"*.PDB\"))\n    rn = rna_name.lower()\n    # exact stem match\n    for f in ref_files:\n        if pathlib.Path(f).stem.lower() == rn:\n            return f\n    # contains match fallback\n    for f in ref_files:\n        if rn in pathlib.Path(f).stem.lower():\n            return f\n    return None\n\ndef iter_rna_residues(structure):\n    \"\"\"\n    Return residues likely to be RNA nucleotides.\n    We treat a residue as RNA-ish if it has P and C4' atoms (common in nucleic acids).\n    \"\"\"\n    residues = []\n    for model in structure:\n        for chain in model:\n            for res in chain:\n                if res.id[0] != \" \":  # skip hetero/water\n                    continue\n                atom_names = {a.get_name() for a in res.get_atoms()}\n                if \"P\" in atom_names and \"C4'\" in atom_names:\n                    residues.append(res)\n    return residues\n\ndef collect_coords(residues, atom_list):\n    \"\"\"\n    For each residue, if all atoms in atom_list exist, collect them in order.\n    Returns: coords (N x 3), index map of residues used.\n    \"\"\"\n    coords = []\n    used_res_idx = []\n    for i, res in enumerate(residues):\n        try:\n            atoms = [res[a] for a in atom_list]\n        except KeyError:\n            continue\n        coords.extend([a.get_coord().astype(np.float64) for a in atoms])\n        used_res_idx.append(i)\n    if len(coords) == 0:\n        return np.zeros((0, 3), dtype=np.float64), []\n    return np.vstack(coords), used_res_idx\n\ndef kabsch_fit(P, Q):\n    \"\"\"\n    Fit P -> Q using Kabsch. P and Q are (N,3).\n    Returns: (R, t, rmsd)\n    \"\"\"\n    assert P.shape == Q.shape and P.shape[1] == 3\n    Pc = P.mean(axis=0)\n    Qc = Q.mean(axis=0)\n    P0 = P - Pc\n    Q0 = Q - Qc\n\n    C = P0.T @ Q0\n    V, S, Wt = np.linalg.svd(C)\n    d = np.sign(np.linalg.det(V @ Wt))\n    D = np.diag([1.0, 1.0, d])\n    R = V @ D @ Wt\n    t = Qc - (Pc @ R)\n    P_aligned = (P @ R) + t\n    rmsd = np.sqrt(np.mean(np.sum((P_aligned - Q) ** 2, axis=1)))\n    return R, t, float(rmsd)\n\ndef apply_transform_to_structure(structure, R, t):\n    \"\"\"Apply rotation R and translation t to all atoms in structure (in-place).\"\"\"\n    for atom in structure.get_atoms():\n        x = atom.get_coord().astype(np.float64)\n        atom.set_coord((x @ R) + t)\n\ndef backbone_rmsd(pred_pdb, ref_pdb, save_aligned_pred_pdb=None):\n    \"\"\"\n    Compute backbone RMSD after fitting pred->ref on BACKBONE atoms.\n    Also returns how many atoms were actually used in the fit.\n    \"\"\"\n    pred = parser.get_structure(\"pred\", pred_pdb)\n    ref  = parser.get_structure(\"ref\",  ref_pdb)\n\n    pred_res = iter_rna_residues(pred)\n    ref_res  = iter_rna_residues(ref)\n\n    # align by residue order; use only overlapping residues\n    n = min(len(pred_res), len(ref_res))\n    pred_res = pred_res[:n]\n    ref_res  = ref_res[:n]\n\n    P, usedP = collect_coords(pred_res, BACKBONE)\n    Q, usedQ = collect_coords(ref_res,  BACKBONE)\n\n    # usedP and usedQ are residue indices within sliced lists; we need intersection positions\n    used = sorted(set(usedP).intersection(set(usedQ)))\n    if len(used) == 0:\n        return None\n\n    # rebuild coords only for residues in intersection, preserving residue order\n    P_list, Q_list = [], []\n    for i in used:\n        resP = pred_res[i]\n        resQ = ref_res[i]\n        # ensure atoms exist in both\n        try:\n            atomsP = [resP[a] for a in BACKBONE]\n            atomsQ = [resQ[a] for a in BACKBONE]\n        except KeyError:\n            continue\n        P_list.extend([a.get_coord().astype(np.float64) for a in atomsP])\n        Q_list.extend([a.get_coord().astype(np.float64) for a in atomsQ])\n\n    if len(P_list) == 0:\n        return None\n\n    P = np.vstack(P_list)\n    Q = np.vstack(Q_list)\n\n    R, t, rmsd = kabsch_fit(P, Q)\n\n    # save aligned predicted pdb if requested\n    if save_aligned_pred_pdb:\n        pred2 = parser.get_structure(\"pred2\", pred_pdb)  # reload so we don’t mutate other computations\n        apply_transform_to_structure(pred2, R, t)\n        io = PDBIO()\n        io.set_structure(pred2)\n        io.save(save_aligned_pred_pdb)\n\n    return {\n        \"backbone_rmsd_A\": rmsd,\n        \"fit_atoms\": int(P.shape[0]),\n        \"fit_residues\": int(len(used)),  # residues contributing full backbone\n        \"pred_residues_total\": int(len(iter_rna_residues(parser.get_structure(\"p\", pred_pdb)))),\n        \"ref_residues_total\":  int(len(iter_rna_residues(parser.get_structure(\"r\", ref_pdb)))),\n    }\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T14:49:28.992018Z","iopub.execute_input":"2026-01-13T14:49:28.992518Z","iopub.status.idle":"2026-01-13T14:49:29.009409Z","shell.execute_reply.started":"2026-01-13T14:49:28.992487Z","shell.execute_reply":"2026-01-13T14:49:29.008516Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Batch over predicted folders\nrows = []\npred_folders = sorted([p for p in glob.glob(os.path.join(PRED_ROOT, \"*\")) if os.path.isdir(p)])\n\nfor folder in pred_folders:\n    rna_name = os.path.basename(folder)\n    pred_pdb = os.path.join(folder, \"unrelaxed_model.pdb\")\n    if not os.path.exists(pred_pdb):\n        # try any pdb inside folder\n        anyp = glob.glob(os.path.join(folder, \"*.pdb\"))\n        if anyp:\n            pred_pdb = anyp[0]\n        else:\n            continue\n\n    ref_pdb = find_reference_pdb(rna_name)\n    if ref_pdb is None:\n        rows.append({\n            \"rna\": rna_name,\n            \"status\": \"NO_REFERENCE_MATCH\",\n            \"pred_pdb\": pred_pdb,\n            \"ref_pdb\": None,\n        })\n        continue\n\n    aligned_pred = os.path.join(ALIGNED_OUT, f\"{rna_name}_pred_aligned_to_ref.pdb\")\n\n    try:\n        out = backbone_rmsd(pred_pdb, ref_pdb, save_aligned_pred_pdb=aligned_pred)\n        if out is None:\n            rows.append({\n                \"rna\": rna_name,\n                \"status\": \"FAILED_NO_MATCHED_BACKBONE\",\n                \"pred_pdb\": pred_pdb,\n                \"ref_pdb\": ref_pdb,\n            })\n            continue\n\n        rows.append({\n            \"rna\": rna_name,\n            \"status\": \"OK\",\n            \"pred_pdb\": pred_pdb,\n            \"ref_pdb\": ref_pdb,\n            \"aligned_pred_pdb\": aligned_pred,\n            **out\n        })\n    except Exception as e:\n        rows.append({\n            \"rna\": rna_name,\n            \"status\": f\"ERROR: {type(e).__name__}\",\n            \"error\": str(e),\n            \"pred_pdb\": pred_pdb,\n            \"ref_pdb\": ref_pdb,\n        })\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T14:49:36.470407Z","iopub.execute_input":"2026-01-13T14:49:36.471211Z","iopub.status.idle":"2026-01-13T14:49:37.525026Z","shell.execute_reply.started":"2026-01-13T14:49:36.471178Z","shell.execute_reply":"2026-01-13T14:49:37.524372Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df = pd.DataFrame(rows).sort_values([\"status\", \"rna\"])\ndf.to_csv(CSV_OUT, index=False)\n\nprint(\"Saved CSV:\", CSV_OUT)\nprint(\"Saved aligned PDBs to:\", ALIGNED_OUT)\ndf","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T14:49:54.585956Z","iopub.execute_input":"2026-01-13T14:49:54.586269Z","iopub.status.idle":"2026-01-13T14:49:54.631443Z","shell.execute_reply.started":"2026-01-13T14:49:54.586244Z","shell.execute_reply":"2026-01-13T14:49:54.630563Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import py3Dmol\n\ndef show_pair(ref_pdb, aligned_pred_pdb, width=700, height=450):\n    with open(ref_pdb) as f:\n        ref_txt = f.read()\n    with open(aligned_pred_pdb) as f:\n        pred_txt = f.read()\n\n    view = py3Dmol.view(width=width, height=height)\n    view.addModel(ref_txt, \"pdb\")\n    view.setStyle({\"model\": 0}, {\"cartoon\": {\"color\": \"orange\"}})\n\n    view.addModel(pred_txt, \"pdb\")\n    view.setStyle({\"model\": 1}, {\"cartoon\": {\"color\": \"cyan\"}})\n\n    view.zoomTo()\n    return view\n\n# Show all OK pairs\nok = df[df[\"status\"] == \"OK\"].copy()\nfor _, r in ok.iterrows():\n    print(r[\"rna\"], \"backbone RMSD (Å) =\", r[\"backbone_rmsd_A\"], \"atoms =\", r[\"fit_atoms\"])\n    display(show_pair(r[\"ref_pdb\"], r[\"aligned_pred_pdb\"]))\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import py3Dmol\n\ndef show_pair(ref_pdb, aligned_pred_pdb, width=700, height=450):\n    with open(ref_pdb) as f:\n        ref_txt = f.read()\n    with open(aligned_pred_pdb) as f:\n        pred_txt = f.read()\n\n    view = py3Dmol.view(width=width, height=height)\n    view.addModel(ref_txt, \"pdb\")\n    view.setStyle({\"model\": 0}, {\"cartoon\": {\"color\": \"orange\"}})\n\n    view.addModel(pred_txt, \"pdb\")\n    view.setStyle({\"model\": 1}, {\"cartoon\": {\"color\": \"cyan\"}})\n\n    view.zoomTo()\n    return view\n\n# Show all OK pairs\nok = df[df[\"status\"] == \"OK\"].copy()\nfor _, r in ok.iterrows():\n    print(r[\"rna\"], \"backbone RMSD (Å) =\", r[\"backbone_rmsd_A\"], \"atoms =\", r[\"fit_atoms\"])\n    display(show_pair(r[\"ref_pdb\"], r[\"aligned_pred_pdb\"]))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T14:58:40.194027Z","iopub.execute_input":"2026-01-13T14:58:40.194761Z","iopub.status.idle":"2026-01-13T14:58:40.316747Z","shell.execute_reply.started":"2026-01-13T14:58:40.194726Z","shell.execute_reply":"2026-01-13T14:58:40.315898Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"##  RMSD TM-Score using US-ALIGN\n\nUS-ALIGN Paper : https://www.nature.com/articles/s41592-022-01585-1","metadata":{}},{"cell_type":"code","source":"%%bash\nset -e\n\ncd /kaggle/working\nrm -rf USalign\ngit clone https://github.com/pylelab/USalign.git\ncd USalign\n\nmake\n\n# Put binary on PATH for notebook session\necho \"export PATH=/kaggle/working/USalign:\\$PATH\" >> ~/.bashrc\nexport PATH=/kaggle/working/USalign:$PATH\n\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Add it to PATH inside the notebook session","metadata":{}},{"cell_type":"code","source":"import os\nos.environ[\"PATH\"] = \"/kaggle/working/USalign:\" + os.environ[\"PATH\"]\n!which USalign","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T15:04:02.831063Z","iopub.execute_input":"2026-01-13T15:04:02.831888Z","iopub.status.idle":"2026-01-13T15:04:03.803788Z","shell.execute_reply.started":"2026-01-13T15:04:02.831848Z","shell.execute_reply":"2026-01-13T15:04:03.802909Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"REF_DIR   = \"/kaggle/input/das-14-rnas-benchmarks/Das-14-RNAs\"\nPRED_ROOT = \"/kaggle/working/rhofold_Das_14_outputs\"\nCSV_OUT   = \"/kaggle/working/usalign_rhofold_das14_benchmark.csv\"\n\nparser = PDBParser(QUIET=True)\n\nBACKBONE = [\"P\", \"O5'\", \"C5'\", \"C4'\", \"C3'\", \"O3'\"]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T15:00:13.657313Z","iopub.execute_input":"2026-01-13T15:00:13.657690Z","iopub.status.idle":"2026-01-13T15:00:13.661740Z","shell.execute_reply.started":"2026-01-13T15:00:13.657661Z","shell.execute_reply":"2026-01-13T15:00:13.660950Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os, glob, pathlib, subprocess, re\n\ndef find_reference_pdb(rna_name: str):\n    ref_files = glob.glob(os.path.join(REF_DIR, \"*.pdb\")) + glob.glob(os.path.join(REF_DIR, \"*.PDB\"))\n    rn = rna_name.lower()\n    # exact stem match\n    for f in ref_files:\n        if pathlib.Path(f).stem.lower() == rn:\n            return f\n    # contains fallback (rarely needed)\n    for f in ref_files:\n        if rn in pathlib.Path(f).stem.lower():\n            return f\n    return None\n\ndef run_usalign(ref_pdb, pred_pdb):\n    \"\"\"\n    Run USalign and parse TM-scoreRNA / RMSD / alignment length / coverage when present.\n    US-align auto-recognizes molecule type per paper and program description. :contentReference[oaicite:1]{index=1}\n    \"\"\"\n    cmd = [\"USalign\", ref_pdb, pred_pdb]\n    res = subprocess.run(cmd, capture_output=True, text=True)\n    out = (res.stdout or \"\") + \"\\n\" + (res.stderr or \"\")\n\n    if res.returncode != 0 and (\"TM-score\" not in out and \"TM-scoreRNA\" not in out):\n        return {\"usalign_status\": f\"ERROR(returncode={res.returncode})\", \"usalign_raw\": out[:2000]}\n\n    # Parse TM-scoreRNA (preferred for RNA) or TM-score\n    tm = None\n    # examples vary; we capture first float after TM-scoreRNA or TM-score\n    m = re.search(r\"TM-scoreRNA\\s*=\\s*([0-9]*\\.[0-9]+|[0-9]+)\", out)\n    if m: tm = float(m.group(1))\n    if tm is None:\n        m = re.search(r\"TM-score\\s*=\\s*([0-9]*\\.[0-9]+|[0-9]+)\", out)\n        if m: tm = float(m.group(1))\n\n    # RMSD\n    rmsd = None\n    m = re.search(r\"RMSD\\s*=\\s*([0-9]*\\.[0-9]+|[0-9]+)\", out)\n    if m: rmsd = float(m.group(1))\n\n    # Aligned length (sometimes shown as \"Aligned length=\")\n    ali_len = None\n    m = re.search(r\"Aligned length\\s*=\\s*(\\d+)\", out)\n    if m: ali_len = int(m.group(1))\n\n    # Coverage sometimes appears explicitly; if not, we compute later from ali_len / ref_len\n    cov = None\n    m = re.search(r\"Coverage\\s*=\\s*([0-9]*\\.[0-9]+|[0-9]+)\", out)\n    if m: cov = float(m.group(1))\n\n    return {\n        \"usalign_status\": \"OK\",\n        \"tm_score_rna\": tm,\n        \"usalign_rmsd_A\": rmsd,\n        \"usalign_aligned_len\": ali_len,\n        \"usalign_coverage_reported\": cov,\n        \"usalign_raw\": out[:2000],\n    }\n\ndef iter_rna_residues(structure):\n    residues = []\n    for model in structure:\n        for chain in model:\n            for res in chain:\n                if res.id[0] != \" \":\n                    continue\n                atom_names = {a.get_name() for a in res.get_atoms()}\n                # simple nucleic acid signature\n                if \"P\" in atom_names and \"C4'\" in atom_names:\n                    residues.append(res)\n    return residues\n\ndef kabsch_rmsd_backbone(pred_pdb, ref_pdb):\n    \"\"\"\n   backbone-only RMSD after simple residue-order Kabsch fit.\n    This is NOT a replacement for US-align; it’s a geometry sanity check.\n    \"\"\"\n    pred = parser.get_structure(\"pred\", pred_pdb)\n    ref  = parser.get_structure(\"ref\",  ref_pdb)\n\n    pred_res = iter_rna_residues(pred)\n    ref_res  = iter_rna_residues(ref)\n    n = min(len(pred_res), len(ref_res))\n    pred_res, ref_res = pred_res[:n], ref_res[:n]\n\n    P_list, Q_list = [], []\n    used_res = 0\n    for i in range(n):\n        rP, rQ = pred_res[i], ref_res[i]\n        try:\n            atomsP = [rP[a] for a in BACKBONE]\n            atomsQ = [rQ[a] for a in BACKBONE]\n        except KeyError:\n            continue\n        P_list.extend([a.get_coord().astype(np.float64) for a in atomsP])\n        Q_list.extend([a.get_coord().astype(np.float64) for a in atomsQ])\n        used_res += 1\n\n    if not P_list:\n        return None\n\n    P = np.vstack(P_list)\n    Q = np.vstack(Q_list)\n\n    Pc, Qc = P.mean(0), Q.mean(0)\n    P0, Q0 = P - Pc, Q - Qc\n    C = P0.T @ Q0\n    V, S, Wt = np.linalg.svd(C)\n    d = np.sign(np.linalg.det(V @ Wt))\n    D = np.diag([1.0, 1.0, d])\n    R = V @ D @ Wt\n    t = Qc - (Pc @ R)\n    P_al = (P @ R) + t\n    rmsd = float(np.sqrt(np.mean(np.sum((P_al - Q)**2, axis=1))))\n\n    return {\"kabsch_bb_rmsd_A\": rmsd, \"kabsch_fit_atoms\": int(P.shape[0]), \"kabsch_fit_residues\": used_res}\n\nrows = []\npred_folders = sorted([p for p in glob.glob(os.path.join(PRED_ROOT, \"*\")) if os.path.isdir(p)])\n\nfor folder in pred_folders:\n    rna = os.path.basename(folder)\n    pred_pdb = os.path.join(folder, \"unrelaxed_model.pdb\")\n    if not os.path.exists(pred_pdb):\n        # fallback: any pdb\n        pdbs = glob.glob(os.path.join(folder, \"*.pdb\"))\n        if not pdbs:\n            rows.append({\"rna\": rna, \"status\": \"NO_PRED_PDB\", \"pred_folder\": folder})\n            continue\n        pred_pdb = pdbs[0]\n\n    ref_pdb = find_reference_pdb(rna)\n    if ref_pdb is None:\n        rows.append({\"rna\": rna, \"status\": \"NO_REFERENCE_MATCH\", \"pred_pdb\": pred_pdb})\n        continue\n\n    # US-align metrics (primary)\n    ua = run_usalign(ref_pdb, pred_pdb)\n\n    #  geometric sanity check\n    kb = kabsch_rmsd_backbone(pred_pdb, ref_pdb)\n\n    row = {\n        \"rna\": rna,\n        \"status\": ua.get(\"usalign_status\", \"ERROR\"),\n        \"ref_pdb\": ref_pdb,\n        \"pred_pdb\": pred_pdb,\n        \"tm_score_rna\": ua.get(\"tm_score_rna\"),\n        \"usalign_rmsd_A\": ua.get(\"usalign_rmsd_A\"),\n        \"usalign_aligned_len\": ua.get(\"usalign_aligned_len\"),\n        \"usalign_coverage_reported\": ua.get(\"usalign_coverage_reported\"),\n    }\n    if kb:\n        row.update(kb)\n    rows.append(row)\n\ndf = pd.DataFrame(rows).sort_values([\"status\", \"rna\"])\ndf.to_csv(CSV_OUT, index=False)\n\nCSV_OUT, df\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T15:05:14.200049Z","iopub.execute_input":"2026-01-13T15:05:14.200365Z","iopub.status.idle":"2026-01-13T15:05:14.638923Z","shell.execute_reply.started":"2026-01-13T15:05:14.200340Z","shell.execute_reply":"2026-01-13T15:05:14.638078Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import py3Dmol\nimport numpy as np\nfrom Bio.PDB import PDBParser, PDBIO\n\nparser = PDBParser(QUIET=True)\nBACKBONE = [\"P\", \"O5'\", \"C5'\", \"C4'\", \"C3'\", \"O3'\"]\n\ndef iter_rna_residues(structure):\n    residues = []\n    for model in structure:\n        for chain in model:\n            for res in chain:\n                if res.id[0] != \" \":\n                    continue\n                atom_names = {a.get_name() for a in res.get_atoms()}\n                if \"P\" in atom_names and \"C4'\" in atom_names:\n                    residues.append(res)\n    return residues\n\ndef kabsch_transform(P, Q):\n    Pc, Qc = P.mean(0), Q.mean(0)\n    P0, Q0 = P - Pc, Q - Qc\n    C = P0.T @ Q0\n    V, S, Wt = np.linalg.svd(C)\n    d = np.sign(np.linalg.det(V @ Wt))\n    D = np.diag([1.0, 1.0, d])\n    R = V @ D @ Wt\n    t = Qc - (Pc @ R)\n    return R, t\n\ndef save_pred_aligned_kabsch(pred_pdb, ref_pdb, out_pdb):\n    pred = parser.get_structure(\"pred\", pred_pdb)\n    ref  = parser.get_structure(\"ref\",  ref_pdb)\n\n    pred_res = iter_rna_residues(pred)\n    ref_res  = iter_rna_residues(ref)\n    n = min(len(pred_res), len(ref_res))\n    pred_res, ref_res = pred_res[:n], ref_res[:n]\n\n    P_list, Q_list = [], []\n    for i in range(n):\n        rP, rQ = pred_res[i], ref_res[i]\n        try:\n            atomsP = [rP[a] for a in BACKBONE]\n            atomsQ = [rQ[a] for a in BACKBONE]\n        except KeyError:\n            continue\n        P_list.extend([a.get_coord().astype(np.float64) for a in atomsP])\n        Q_list.extend([a.get_coord().astype(np.float64) for a in atomsQ])\n\n    P = np.vstack(P_list); Q = np.vstack(Q_list)\n    R, t = kabsch_transform(P, Q)\n\n    # apply to all atoms in pred\n    for atom in pred.get_atoms():\n        x = atom.get_coord().astype(np.float64)\n        atom.set_coord((x @ R) + t)\n\n    io = PDBIO()\n    io.set_structure(pred)\n    io.save(out_pdb)\n    return out_pdb\n\ndef show_pair(ref_pdb, pred_aligned_pdb, width=700, height=450):\n    with open(ref_pdb) as f: ref_txt = f.read()\n    with open(pred_aligned_pdb) as f: pred_txt = f.read()\n\n    v = py3Dmol.view(width=width, height=height)\n    v.addModel(ref_txt, \"pdb\")\n    v.setStyle({\"model\": 0}, {\"cartoon\": {\"color\": \"orange\"}})\n    v.addModel(pred_txt, \"pdb\")\n    v.setStyle({\"model\": 1}, {\"cartoon\": {\"color\": \"cyan\"}})\n    v.zoomTo()\n    return v\n\n# first entry\nok = df[df[\"status\"]==\"OK\"].iloc[0]\nrna = ok[\"rna\"]\naligned_path = f\"/kaggle/working/{rna}_pred_kabsch_aligned.pdb\"\nsave_pred_aligned_kabsch(ok[\"pred_pdb\"], ok[\"ref_pdb\"], aligned_path)\n\nprint(\"RNA:\", rna)\nprint(\"US-align TM-scoreRNA/TM:\", ok[\"tm_score_rna\"], \"US-align RMSD (Å):\", ok[\"usalign_rmsd_A\"])\ndisplay(show_pair(ok[\"ref_pdb\"], aligned_path))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-01-13T15:06:45.251927Z","iopub.execute_input":"2026-01-13T15:06:45.252609Z","iopub.status.idle":"2026-01-13T15:06:45.302211Z","shell.execute_reply.started":"2026-01-13T15:06:45.252579Z","shell.execute_reply":"2026-01-13T15:06:45.301392Z"}},"outputs":[],"execution_count":null}]}