{"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":"gpu","dataSources":[{"sourceId":87793,"databundleVersionId":11403143,"sourceType":"competition"},{"sourceId":10855324,"sourceType":"datasetVersion","datasetId":6742586},{"sourceId":11031941,"sourceType":"datasetVersion","datasetId":6869784},{"sourceId":224830487,"sourceType":"kernelVersion"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"we are using commandline \"python inference.py\" to verify correct installtion in this demo. You can instead use rhofold+ net api and write your python code for faster processing.\n\n1. setup rhofold+\nplease follow instruction at https://github.com/ml4bio/RhoFold\n   \n2. setup usalign\nplease follow instruction at https://github.com/pylelab/USalign  \nuse:  g++ -static -O3 -ffast-math -lm -o USalign USalign.cpp  ","metadata":{}},{"cell_type":"code","source":"try:\n    import Bio\nexcept:\n    #for rhofold+ #####################\n    !pip install biopython\n    !pip install ml-collections\n    !pip install python-box\n    !pip install dm-tree\n    !pip install openmm[cuda12]\n\n\n\n\n\nfrom copy import deepcopy\n\nimport pandas as pd\nfrom Bio.PDB import Atom, Model, Chain, Residue, Structure, PDBParser\nfrom Bio import SeqIO\nimport os, sys\nimport re\nimport numpy as np\n\nimport matplotlib\nimport matplotlib.pyplot as plt\n\nprint('IMPORT OK !!!!')","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-03-14T16:57:02.849052Z","iopub.execute_input":"2025-03-14T16:57:02.849358Z","iopub.status.idle":"2025-03-14T16:57:02.858596Z","shell.execute_reply.started":"2025-03-14T16:57:02.849339Z","shell.execute_reply":"2025-03-14T16:57:02.857713Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"PYTHON = sys.executable\nprint('PYTHON',PYTHON)\n\nRHONET_DIR=\\\n'/kaggle/input/data-for-demo-for-rhofold-plus-with-kaggle-msa/RhoFold-main'\n#'<your downloaded rhofold repo>/RhoFold-main'\n\nUSALIGN = \\\n'/kaggle/working//USalign'\n#'<your us align path>/USalign'\n\nos.system('cp /kaggle/input/usalign/USalign /kaggle/working/')\nos.system('sudo chmod u+x /kaggle/working//USalign')\n\n\nDATA_KAGGLE_DIR = '/kaggle/input/stanford-rna-3d-folding'\nSEQ_DF = pd.read_csv(f'{DATA_KAGGLE_DIR}/train_sequences.csv')\nLABEL_DF = pd.read_csv(f'{DATA_KAGGLE_DIR}/train_labels.csv')\nLABEL_DF['target_id'] = LABEL_DF['ID'].apply(lambda x: '_'.join(x.split('_')[:-1]))\n\n\n# helper ----\nclass dotdict(dict):\n\t__setattr__ = dict.__setitem__\n\t__delattr__ = dict.__delitem__\n\n\tdef __getattr__(self, name):\n\t\ttry:\n\t\t\treturn self[name]\n\t\texcept KeyError:\n\t\t\traise AttributeError(name)\n\n# visualisation helper ----\ndef set_aspect_equal(ax):\n\tx_limits = ax.get_xlim()\n\ty_limits = ax.get_ylim()\n\tz_limits = ax.get_zlim()\n\n\t# Compute the mean of each axis\n\tx_middle = np.mean(x_limits)\n\ty_middle = np.mean(y_limits)\n\tz_middle = np.mean(z_limits)\n\n\t# Compute the max range across all axes\n\tmax_range = max(x_limits[1] - x_limits[0],\n\t\t\t\t\ty_limits[1] - y_limits[0],\n\t\t\t\t\tz_limits[1] - z_limits[0]) / 2.0\n\n\t# Set the new limits to ensure equal scaling\n\tax.set_xlim(x_middle - max_range, x_middle + max_range)\n\tax.set_ylim(y_middle - max_range, y_middle + max_range)\n\tax.set_zlim(z_middle - max_range, z_middle + max_range)\n\n\n\n\n# xyz df helper --------------------\ndef get_truth_df(target_id):\n    truth_df = LABEL_DF[LABEL_DF['target_id'] == target_id]\n    truth_df = truth_df.reset_index(drop=True)\n    return truth_df\n\ndef parse_pdb_to_df(pdb_file, target_id):\n    parser = PDBParser()\n    structure = parser.get_structure('', pdb_file)\n\n    df = []\n    for model in structure:\n        for chain in model:\n            print(chain)\n            chain_data = []\n            for residue in chain:\n                # print(residue)\n                if residue.get_resname() in ['A', 'U', 'G', 'C']:\n                    # Check if the residue has a C1' atom\n                    if 'C1\\'' in residue:\n                        atom = residue['C1\\'']\n                        xyz = atom.get_coord()\n                        resname = residue.get_resname()\n                        resid = residue.get_id()[1]\n\n                        #todo detect discontinous: resid = prev_resid+1\n                        #ID\tresname\tresid\tx_1\ty_1\tz_1\n                        chain_data.append(dict(\n                            ID = target_id+'_'+str(resid),\n                            resname=resname,\n                            resid=resid,\n                            x_1=xyz[0],\n                            y_1=xyz[1],\n                            z_1=xyz[2],\n                        ))\n                        ##print(f\"Residue {resname} {resid}, Atom: {atom.get_name()}, xyz: {xyz}\")\n\n            if len(chain_data)!=0:\n                chain_df = pd.DataFrame(chain_data)\n                df.append(chain_df)\n                ##print(chain_df)\n    return df\n\n# usalign helper --------------------\ndef write_target_line(\n    atom_name, atom_serial, residue_name, chain_id, residue_num, x_coord, y_coord, z_coord, occupancy=1.0, b_factor=0.0, atom_type='P'\n):\n    \"\"\"\n    Writes a single line of PDB format based on provided atom information.\n\n    Args:\n        atom_name (str): Name of the atom (e.g., \"N\", \"CA\").\n        atom_serial (int): Atom serial number.\n        residue_name (str): Residue name (e.g., \"ALA\").\n        chain_id (str): Chain identifier.\n        residue_num (int): Residue number.\n        x_coord (float): X coordinate.\n        y_coord (float): Y coordinate.\n        z_coord (float): Z coordinate.\n        occupancy (float, optional): Occupancy value (default: 1.0).\n        b_factor (float, optional): B-factor value (default: 0.0).\n\n    Returns:\n        str: A single line of PDB string.\n    \"\"\"\n    return f'ATOM  {atom_serial:>5d}  {atom_name:<5s} {residue_name:<3s} {residue_num:>3d}    {x_coord:>8.3f}{y_coord:>8.3f}{z_coord:>8.3f}{occupancy:>6.2f}{b_factor:>6.2f}           {atom_type}\\n'\n\ndef write_xyz_to_pdb(df, pdb_file, xyz_id = 1):\n    resolved_cnt = 0\n    with open(pdb_file, 'w') as target_file:\n        for _, row in df.iterrows():\n            x_coord = row[f'x_{xyz_id}']\n            y_coord = row[f'y_{xyz_id}']\n            z_coord = row[f'z_{xyz_id}']\n\n            if x_coord > -1e17 and y_coord > -1e17 and z_coord > -1e17:\n                resolved_cnt += 1\n                target_line = write_target_line(\n                    atom_name=\"C1'\",\n                    atom_serial=int(row['resid']),\n                    residue_name=row['resname'],\n                    chain_id='0',\n                    residue_num=int(row['resid']),\n                    x_coord=x_coord,\n                    y_coord=y_coord,\n                    z_coord=z_coord,\n                    atom_type='C',\n                )\n                target_file.write(target_line)\n    return resolved_cnt\n\ndef parse_usalign_for_tm_score(output):\n    # Extract TM-score based on length of reference structure (second)\n    tm_score_match = re.findall(r'TM-score=\\s+([\\d.]+)', output)[1]\n    if not tm_score_match:\n        raise ValueError('No TM score found')\n    return float(tm_score_match)\n\ndef parse_usalign_for_transform(output):\n    # Locate the rotation matrix section\n    matrix_lines = []\n    found_matrix = False\n\n    for line in output.splitlines():\n        if \"The rotation matrix to rotate Structure_1 to Structure_2\" in line:\n            found_matrix = True\n        elif found_matrix and re.match(r'^\\d+\\s+[-\\d.]+\\s+[-\\d.]+\\s+[-\\d.]+\\s+[-\\d.]+$', line):\n            matrix_lines.append(line)\n        elif found_matrix and not line.strip():\n            break  # Stop parsing if an empty line is encountered after the matrix\n\n    # Parse the rotation matrix values\n    rotation_matrix = []\n    for line in matrix_lines:\n        parts = line.split()\n        row_values = list(map(float, parts[1:]))  # Skip the first column (index)\n        rotation_matrix.append(row_values)\n\n    return np.array(rotation_matrix)\n\ndef call_usalign(predict_df, truth_df, verbose=1):\n    truth_pdb = '~truth.pdb'\n    predict_pdb = '~predict.pdb'\n    write_xyz_to_pdb(predict_df, predict_pdb, xyz_id=1)\n    write_xyz_to_pdb(truth_df, truth_pdb, xyz_id=1)\n\n    command = f'{USALIGN} {predict_pdb} {truth_pdb} -atom \" C1\\'\" -m -'\n    output = os.popen(command).read()\n    if verbose==1:\n        print(output)\n    tm_score = parse_usalign_for_tm_score(output)\n    transform = parse_usalign_for_transform(output)\n    return tm_score, transform\n\n\n# msa helper --------------------\ndef read_msa(msa_file):\n    f = open(msa_file, 'r')\n    line = f.readlines()\n\n    msa = []\n    for i in range(0, len(line),2):\n        m = dotdict(\n            comment =line[i],\n            seqence =line[i+1],\n        )\n        assert(m.comment[0]=='>')\n        msa.append(m)\n    return msa\n\n\ndef write_msa(msa_file, msa):\n    line=[]\n    for m in msa:\n        line .append(m.comment)\n        line .append(m.seqence)\n\n    f = open(msa_file, 'wt')\n    f.writelines(line)\n    return msa\n \ndef msa_to_rhonet_file(msa_file, num_msa=5, out_dir='',target_id='xxx'):\n    msa = read_msa(msa_file)\n    msa0 = deepcopy(msa[0])\n    msa0.comment =f'>{target_id}\\n'\n    msa0 = [msa0]\n\n    a3m_file = f'{out_dir}/{target_id}.a3m'\n    fasta_file = f'{out_dir}/{target_id}.fasta'\n    os.makedirs(out_dir, exist_ok=True)\n\n    write_msa(fasta_file, msa0)\n    write_msa(a3m_file, msa[:num_msa])\n\nprint('HELPER OK!!!')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-14T16:57:02.868848Z","iopub.execute_input":"2025-03-14T16:57:02.869053Z","iopub.status.idle":"2025-03-14T16:57:03.156891Z","shell.execute_reply.started":"2025-03-14T16:57:02.869035Z","shell.execute_reply":"2025-03-14T16:57:03.156019Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#start here!!!\n\n\nout_dir   ='/kaggle/working/'\ntarget_id = '1EIY_C'.upper()\nsequence  = 'GCCGAGGUAGCUCAGUUGGUAGAGCAUGCGACUGAAAAUCGCAGUGUCCGCGGUUCGAUUCCGCGCCUCGGCACCA'\nprint('len(sequence):',len(sequence))\n\n\n#1. prepare input\nmsa_file = f'{DATA_KAGGLE_DIR}/MSA/1EIY_C.MSA.fasta'\nmsa_to_rhonet_file(msa_file, num_msa=5, out_dir=out_dir,target_id=target_id)\n\ncmd1 = f'cd {RHONET_DIR}'\n#cmd2 = f'export LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libstdc++.so.6' #optional if you have lib error\ncmd3 = f'{PYTHON} inference.py --input_fas {out_dir}/{target_id}.fasta --input_a3m {out_dir}/{target_id}.a3m --output_dir {out_dir}/ --ckpt ./pretrained/model_20221010_params.pt'\n\n#follow rhofold repo, we use the cmdline:\n#'python inference.py --input_fas ./example/input/3owzA/3owzA.fasta --input_a3m ./example/input/3owzA/3owzA.a3m --output_dir ./example/output/3owzA/ --ckpt ./pretrained/model_20221010_params.pt'\n\n#do inference here!\n#output = os.popen(cmd1+';'+cmd2+';'+cmd3).read()\noutput = os.popen(cmd1+';'+cmd3).read()\nprint(output)\n\n#copy file from local results\n#local_result = '/kaggle/input/data-for-demo-for-rhofold-plus-with-kaggle-msa/rhofold_input_output/.'\n#!cp -a $local_result $out_dir\n\n#expected ouput from rtx A6000 (non ada)\n'''\n2025-03-14 20:40:50,182 - INFO: Constructing RhoFold\n2025-03-14 20:40:51,221 - INFO:     loading ./pretrained/model_20221010_params.pt\n2025-03-14 20:40:51,743 - INFO: Input_fas /media/hp/c30d34ed-0d55-4077-82dc-b56cd13dd548/2025/kaggle/stanford-rna-3d-folding/result/rhofold_00/1EIY_C.fasta\n2025-03-14 20:40:51,743 - INFO: Input_a3m /media/hp/c30d34ed-0d55-4077-82dc-b56cd13dd548/2025/kaggle/stanford-rna-3d-folding/result/rhofold_00/1EIY_C.a3m\n2025-03-14 20:40:51,743 - INFO: Started RhoFold Inference\n2025-03-14 20:40:51,755 - INFO:     Inference using device cuda\n2025-03-14 20:40:54,523 - INFO:     Export PDB file to /media/hp/c30d34ed-0d55-4077-82dc-b56cd13dd548/2025/kaggle/stanford-rna-3d-folding/result/rhofold_00//unrelaxed_model.pdb\n2025-03-14 20:40:54,523 - INFO: Finished RhoFold Inference in 2.780 seconds\n2025-03-14 20:40:54,523 - INFO: Started Amber Relaxation : 1000 iterations\n2025-03-14 20:40:54,523 - INFO:     AmberRelaxation: Using OpenCL\n2025-03-14 20:41:09,410 - INFO:     Minimizing ...\n2025-03-14 20:42:38,203 - INFO:     Energy at Minima is -505932.780 kcal/mol\n2025-03-14 20:42:38,362 - INFO:     Export PDB file to /media/hp/c30d34ed-0d55-4077-82dc-b56cd13dd548/2025/kaggle/stanford-rna-3d-folding/result/rhofold_00//relaxed_1000_model.pdb\n2025-03-14 20:42:38,363 - INFO: Finished Amber Relaxation : 1000 iterations in 103.840 seconds\n'''\n\n#expected ouput from P100\n'''\nlen(sequence): 76\n2025-03-14 16:03:33,110 - INFO: Constructing RhoFold\n2025-03-14 16:03:34,429 - INFO:     loading ./pretrained/model_20221010_params.pt\n2025-03-14 16:03:35,088 - INFO: Input_fas /kaggle/working//1EIY_C.fasta\n2025-03-14 16:03:35,089 - INFO: Input_a3m /kaggle/working//1EIY_C.a3m\n2025-03-14 16:03:35,089 - INFO: Started RhoFold Inference\n2025-03-14 16:03:35,093 - INFO:     Inference using device cuda\n2025-03-14 16:03:40,177 - INFO:     Export PDB file to /kaggle/working///unrelaxed_model.pdb\n2025-03-14 16:03:40,177 - INFO: Finished RhoFold Inference in 5.088 seconds\n2025-03-14 16:03:40,177 - INFO: Started Amber Relaxation : 1000 iterations\n2025-03-14 16:03:40,177 - INFO:     AmberRelaxation: Using OpenCL\n2025-03-14 16:04:03,056 - INFO:     Minimizing ...\n2025-03-14 16:08:49,875 - INFO:     Energy at Minima is -497896.334 kcal/mol\n2025-03-14 16:08:50,044 - INFO:     Export PDB file to /kaggle/working///relaxed_1000_model.pdb\n2025-03-14 16:08:50,047 - INFO: Finished Amber Relaxation : 1000 iterations in 309.869 seconds\n\n'''","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-14T17:10:39.450541Z","iopub.execute_input":"2025-03-14T17:10:39.450917Z","iopub.status.idle":"2025-03-14T17:17:45.743142Z","shell.execute_reply.started":"2025-03-14T17:10:39.450881Z","shell.execute_reply":"2025-03-14T17:17:45.742380Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"#visualise prediction and compute tm score\n\npredict_relax_df = parse_pdb_to_df(f'{out_dir}/relaxed_1000_model.pdb', target_id)\npredict_unrelax_df = parse_pdb_to_df(f'{out_dir}/unrelaxed_model.pdb', target_id)\n\nassert len(predict_relax_df)==1\nassert len(predict_unrelax_df)==1\npredict_relax_df = predict_relax_df[0]\npredict_unrelax_df = predict_unrelax_df[0]\n\nprint(predict_relax_df)\nprint(predict_unrelax_df)\n\ntruth_df = get_truth_df(target_id)\nprint(truth_df)\n\ntm_score_relax, transform_relax = call_usalign(predict_relax_df, truth_df, verbose=1)\ntm_score_unrelax, transform_unrelax= call_usalign(predict_unrelax_df, truth_df, verbose=0)\n\nprint('tm_score_relax', tm_score_relax)\nprint('tm_score_unrelax', tm_score_unrelax)\nprint('transform_relax\\n', transform_relax)\nprint('transform_unrelax\\n', transform_unrelax)\nzz=0\n\nif 1:\n    COLOR = ['red', 'blue', 'green', 'black', 'yellow', 'cyan', 'magenta']\n    fig = plt.figure(figsize=(10, 10))\n    ax = fig.add_subplot(111, projection='3d')\n    # ax.clear()\n\n    #unrelax\n    coord = predict_unrelax_df[['x_1', 'y_1', 'z_1']].to_numpy().astype('float32')\n    coord = coord@transform_unrelax[:,1:].T + transform_unrelax[:,[0]].T\n    x, y, z = coord[:, 0], coord[:, 1], coord[:, 2]\n    ax.scatter(x, y, z, c='red', s=30, alpha=1)\n    ax.plot(x, y, z, color='red', linewidth=1, alpha=1, label=f'unrelax (tm:{tm_score_unrelax:0.3f})')\n\n\n    #relax\n    coord = predict_relax_df[['x_1', 'y_1', 'z_1']].to_numpy().astype('float32')\n    coord = coord@transform_relax[:,1:].T + transform_relax[:,[0]].T\n    x, y, z = coord[:, 0], coord[:, 1], coord[:, 2]\n    ax.scatter(x, y, z, c='orange', s=30, alpha=1)\n    ax.plot(x, y, z, color='orange', linewidth=1, alpha=1, label=f'relax (tm:{tm_score_relax:0.3f})')\n\n    # truth\n    truth = truth_df[['x_1', 'y_1', 'z_1']].to_numpy().astype('float32')\n    x, y, z = truth[:, 0], truth[:, 1], truth[:, 2]\n    ax.scatter(x, y, z, c='black', s=30, alpha=1)\n    ax.plot(x, y, z, color='black', linewidth=1, alpha=1, label=f'truth')\n\n    set_aspect_equal(ax)\n    plt.legend()\n    plt.show()\n    # plt.waitforbuttonpress()\n    plt.close()\n    ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-03-14T17:19:49.334046Z","iopub.execute_input":"2025-03-14T17:19:49.334361Z","iopub.status.idle":"2025-03-14T17:19:49.667829Z","shell.execute_reply.started":"2025-03-14T17:19:49.334341Z","shell.execute_reply":"2025-03-14T17:19:49.667003Z"}},"outputs":[],"execution_count":null}]}