{"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":"none","dataSources":[{"sourceId":87793,"databundleVersionId":11553390,"sourceType":"competition"},{"sourceId":10855324,"sourceType":"datasetVersion","datasetId":6742586},{"sourceId":232150220,"sourceType":"kernelVersion"},{"sourceId":232961333,"sourceType":"kernelVersion"}],"dockerImageVersionId":30918,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"## What this notebook do\nIn this notebook series, I visualize predictions&truth for some test sequences. Truth was obtained from train dataset(train&test have identical sequences(see this notebook:https://www.kaggle.com/code/nanacat0520/stanfordrna-train-test-have-identical-sequence )).","metadata":{}},{"cell_type":"markdown","source":"## How to use this notebook¶¶\nYou can put your submission.csv in input and run this notebook to visualize your submission. Also, check ohter notebooks(https://www.kaggle.com/code/nanacat0520/stanford-rna-submission-to-visualization-protenix, https://www.kaggle.com/code/nanacat0520/stanford-rna-submission-to-vis-rho-ribonanza for visualization of other public code submissions)","metadata":{}},{"cell_type":"markdown","source":"Credit: @fernandosr85 for inference notebook","metadata":{}},{"cell_type":"markdown","source":"## Insights obtained from this notebook¶\n\n・gnn sometimes show higher score than protenix","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\nimport pandas as pd\n\nimport matplotlib\nimport matplotlib.pyplot as plt\n!pip install kagglehub\nimport kagglehub\n\nprint('IMPORT OK !!!!')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T03:45:30.771397Z","iopub.execute_input":"2025-04-12T03:45:30.771793Z","iopub.status.idle":"2025-04-12T03:46:10.316454Z","shell.execute_reply.started":"2025-04-12T03:45:30.771765Z","shell.execute_reply":"2025-04-12T03:46:10.315036Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"PYTHON = sys.executable\nprint('PYTHON',PYTHON)\n\nusalign_path = kagglehub.dataset_download('metric/usalign')\nUSALIGN = \\\n'/kaggle/working//USalign'\n#'<your us align path>/USalign'\n\nos.system('cp /kaggle/input/usalign/USalign /kaggle/working/')\nos.system(' chmod u+x /kaggle/working//USalign')\n\n\nsubmission = pd.read_csv('/kaggle/input/rna-3d-fold-hybrid-template-nn-structure/submission.csv')\nLABEL_DF = pd.read_csv('/kaggle/input/stanford-rna-3d-folding/train_labels.csv')\n\nLABEL_DF[\"pdb_id\"] = LABEL_DF[\"ID\"].apply(lambda x: x.split(\"_\")[0]+'_'+x.split(\"_\")[1])\nsubmission['submission_id'] = submission['ID'].str.split('_').str[0]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T03:46:10.318176Z","iopub.execute_input":"2025-04-12T03:46:10.319003Z","iopub.status.idle":"2025-04-12T03:46:11.047674Z","shell.execute_reply.started":"2025-04-12T03:46:10.318965Z","shell.execute_reply":"2025-04-12T03:46:11.046667Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"class 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\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                if len(chain_data)!=0:\n                    chain_df = pd.DataFrame(chain_data)\n                    df.append(chain_df)\n    return df\n\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\n\n\ndef write_xyz_to_pdb(df, pdb_file, xyz_id ):\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,xyz_id, verbose=1):\n    truth_pdb = '~truth.pdb'\n    predict_pdb = '~predict.pdb'\n    write_xyz_to_pdb(predict_df, predict_pdb, xyz_id)\n    write_xyz_to_pdb(truth_df, truth_pdb, 1)\n\n    command = f'{USALIGN} {predict_pdb} {truth_pdb} -atom \"C1\\'\" -m -'\n    output = os.popen(command).read()\n\n    if verbose == 1:\n        print(\"=== USAlign Output ===\")\n        print(output)\n\n        matches = re.findall(r'TM-score=\\s+([\\d.]+)', output)\n        print(\"=== TM-score Matches ===\")\n        print(matches)\n\n    tm_score = parse_usalign_for_tm_score(output)\n    transform = parse_usalign_for_transform(output)\n    return tm_score, transform\n\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T03:46:11.049464Z","iopub.execute_input":"2025-04-12T03:46:11.049818Z","iopub.status.idle":"2025-04-12T03:46:11.066162Z","shell.execute_reply.started":"2025-04-12T03:46:11.049790Z","shell.execute_reply":"2025-04-12T03:46:11.064860Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def get_truth_df(target_id):\n    truth_df = LABEL_DF[LABEL_DF['pdb_id'] == target_id]\n    truth_df = truth_df.reset_index(drop=True)\n    return truth_df\n\nsubmission_truth_dict = {\n    \"R1126\": \"8TVZ_C\",\n    \"R1128\": \"8BTZ_A\",\n    \"R1138\": \"7PTL_B\",\n    \"R1149\": \"8UYS_A\",\n    \"R1156\": \"8UYJ_A\",\n    \"R1189\": \"7YR7_A\"\n}","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T03:46:11.067930Z","iopub.execute_input":"2025-04-12T03:46:11.068337Z","iopub.status.idle":"2025-04-12T03:46:11.103296Z","shell.execute_reply.started":"2025-04-12T03:46:11.068303Z","shell.execute_reply":"2025-04-12T03:46:11.102185Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 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)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T03:46:11.104408Z","iopub.execute_input":"2025-04-12T03:46:11.104804Z","iopub.status.idle":"2025-04-12T03:46:11.122140Z","shell.execute_reply.started":"2025-04-12T03:46:11.104767Z","shell.execute_reply":"2025-04-12T03:46:11.120855Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def submission_to_visual_2(submission):\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    prediction_count = (submission.shape[1]-4)/3\n    for key, value in submission_truth_dict.items():\n        pred_sorted_sequence = submission[submission['submission_id'] == key]\n        pred_sorted_sequence = pred_sorted_sequence.reset_index(drop=True)\n        truth_df = get_truth_df(value)\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        for i in range(1, int(prediction_count)+1):\n            xyz_id = i\n            fig = plt.figure(figsize=(8, 8))\n            ax = fig.add_subplot(111, projection='3d')\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='truth')\n            coord_tmp = pred_sorted_sequence[[f'x_{i}', f'y_{i}', f'z_{i}']].to_numpy().astype('float32')\n            pred_sorted_sequence_tmp = pred_sorted_sequence[['ID', 'resname', 'resid', f'x_{i}', f'y_{i}', f'z_{i}']].copy()\n            pred_sorted_sequence_tmp.rename(columns={\n                f'x_{i}': 'x_1',\n                f'y_{i}': 'y_1',\n                f'z_{i}': 'z_1',\n            }, inplace=True)\n            tm_score_submission, transform_submission = call_usalign(pred_sorted_sequence_tmp, truth_df, xyz_id=1, verbose=0)\n            coord_tmp = coord_tmp @ transform_submission[:, 1:].T + transform_submission[:, [0]].T\n            x, y, z = coord_tmp[:, 0], coord_tmp[:, 1], coord_tmp[:, 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'prediction {i} (TM: {tm_score_submission:.3f})')\n            set_aspect_equal(ax)\n            ax.set_title(f'Prediction {i} vs Truth ({value})')\n            plt.legend()\n            plt.show()\n            ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T03:46:11.123364Z","iopub.execute_input":"2025-04-12T03:46:11.123691Z","iopub.status.idle":"2025-04-12T03:46:11.146413Z","shell.execute_reply.started":"2025-04-12T03:46:11.123657Z","shell.execute_reply":"2025-04-12T03:46:11.145227Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_to_visual_2(submission)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-04-12T03:46:11.147603Z","iopub.execute_input":"2025-04-12T03:46:11.147994Z","iopub.status.idle":"2025-04-12T03:46:32.814665Z","shell.execute_reply.started":"2025-04-12T03:46:11.147958Z","shell.execute_reply":"2025-04-12T03:46:32.813160Z"}},"outputs":[],"execution_count":null}]}