Compare TorRNA and RNA Puzzles Submissions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import numpy as np\n",
    "import math\n",
    "from tqdm import tqdm\n",
    "import wandb\n",
    "import pickle\n",
    "import pandas as pd\n",
    "\n",
    "from data.rna_object import RNA\n",
    "from data.dataloader import get_train_test_dataloaders, RNATorsionalAnglesDataset, collate_fn\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "\n",
    "from models.transformer import TorsionalAnglesTransformerDecoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found precomputed test dataset\n",
      "Found precomputed train dataset\n",
      "Found RNA objects precomputed\n",
      "Found precomputed FASTA sequences\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(0)\n",
    "np.random.seed(0)\n",
    "\n",
    "pdbs_path=\"/home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs/\"\n",
    "processed_dir = \"/ssd_scratch/users/sriram.devata/rna_structure/dataset/\"\n",
    "perfect_pdb_files_train_val_test_path=\"/home2/sriram.devata/rna_project/cdhit/torrna_train_val_test.pkl\"\n",
    "\n",
    "with open(perfect_pdb_files_train_val_test_path, \"rb\") as fp:\n",
    "        training_pdbs, validation_pdbs, testing_pdbs = pickle.load(fp)\n",
    "training_files = list()\n",
    "for each_training_pdb in training_pdbs:\n",
    "        training_files.append(f\"{pdbs_path}/{each_training_pdb}.pdb\")\n",
    "testing_files = list()\n",
    "for each_testing_pdb in testing_pdbs:\n",
    "        testing_files.append(f\"{pdbs_path}/{each_testing_pdb}.pdb\")\n",
    "\n",
    "pdbs_to_predict = 1000\n",
    "list_of_pdbs_to_predict = testing_files[:pdbs_to_predict]\n",
    "predict_dataset = RNATorsionalAnglesDataset(list_of_pdbs_to_predict, processed_dir=processed_dir, type_dataset=\"test\")\n",
    "predict_dataloader = DataLoader(predict_dataset, collate_fn=collate_fn, batch_size=32)\n",
    "\n",
    "train_dataset = RNATorsionalAnglesDataset(training_files, processed_dir=processed_dir, type_dataset=\"train\")\n",
    "train_dataloader = DataLoader(train_dataset, collate_fn=collate_fn, batch_size=32)\n",
    "\n",
    "\n",
    "if os.path.exists(f\"/ssd_scratch/users/sriram.devata/rna_objects_temp.pkl\"):\n",
    "        with open(f\"/ssd_scratch/users/sriram.devata/rna_objects_temp.pkl\", \"rb\") as input_file:\n",
    "                print(f\"Found RNA objects precomputed\")\n",
    "                rna_objects = pickle.load(input_file)\n",
    "else:\n",
    "        rna_objects = list()\n",
    "        print(f\"Making RNA objects\")\n",
    "        for each_pdb_file in tqdm(list_of_pdbs_to_predict):\n",
    "                try:\n",
    "                        rna_objects.append(RNA(each_pdb_file, calc_rna_fm_embeddings=True, load_dssr_dihedrals=True, load_coords=False))\n",
    "                except:\n",
    "                        pass\n",
    "        with open(f\"/ssd_scratch/users/sriram.devata/rna_objects_temp.pkl\", \"wb\") as output_file:\n",
    "                pickle.dump(rna_objects, output_file)\n",
    "        \n",
    "if os.path.exists(f\"predict_seq.fasta\"):\n",
    "        print(f\"Found precomputed FASTA sequences\")\n",
    "else:\n",
    "        print(f\"Computing FASTA sequences for SPOT-RNA-1D\")\n",
    "        with open(f\"predict_seq.fasta\", \"w\") as fp:\n",
    "                for ran_object_idx,each_rna_object in enumerate(rna_objects):\n",
    "                        pdb_code = each_rna_object.pdb_path.split('/')[-1].replace(\".pdb\", \"\")\n",
    "                        fp.write(f\">{pdb_code}\\n\")\n",
    "                        fp.write(f\"{each_rna_object.dssr_full_seq}\\n\")\n",
    "        print(f\"Run SPOT-RNA-1D\")\n",
    "        exit(1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "TorsionalAnglesTransformerDecoder(\n",
      "  (transformer_decoder): TransformerDecoder(\n",
      "    (layers): ModuleList(\n",
      "      (0): TransformerDecoderLayer(\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (out_proj): NonDynamicallyQuantizableLinear(in_features=640, out_features=640, bias=True)\n",
      "        )\n",
      "        (multihead_attn): MultiheadAttention(\n",
      "          (out_proj): NonDynamicallyQuantizableLinear(in_features=640, out_features=640, bias=True)\n",
      "        )\n",
      "        (linear1): Linear(in_features=640, out_features=2048, bias=True)\n",
      "        (dropout): Dropout(p=0.2, inplace=False)\n",
      "        (linear2): Linear(in_features=2048, out_features=640, bias=True)\n",
      "        (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (dropout1): Dropout(p=0.2, inplace=False)\n",
      "        (dropout2): Dropout(p=0.2, inplace=False)\n",
      "        (dropout3): Dropout(p=0.2, inplace=False)\n",
      "      )\n",
      "      (1): TransformerDecoderLayer(\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (out_proj): NonDynamicallyQuantizableLinear(in_features=640, out_features=640, bias=True)\n",
      "        )\n",
      "        (multihead_attn): MultiheadAttention(\n",
      "          (out_proj): NonDynamicallyQuantizableLinear(in_features=640, out_features=640, bias=True)\n",
      "        )\n",
      "        (linear1): Linear(in_features=640, out_features=2048, bias=True)\n",
      "        (dropout): Dropout(p=0.2, inplace=False)\n",
      "        (linear2): Linear(in_features=2048, out_features=640, bias=True)\n",
      "        (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (dropout1): Dropout(p=0.2, inplace=False)\n",
      "        (dropout2): Dropout(p=0.2, inplace=False)\n",
      "        (dropout3): Dropout(p=0.2, inplace=False)\n",
      "      )\n",
      "      (2): TransformerDecoderLayer(\n",
      "        (self_attn): MultiheadAttention(\n",
      "          (out_proj): NonDynamicallyQuantizableLinear(in_features=640, out_features=640, bias=True)\n",
      "        )\n",
      "        (multihead_attn): MultiheadAttention(\n",
      "          (out_proj): NonDynamicallyQuantizableLinear(in_features=640, out_features=640, bias=True)\n",
      "        )\n",
      "        (linear1): Linear(in_features=640, out_features=2048, bias=True)\n",
      "        (dropout): Dropout(p=0.2, inplace=False)\n",
      "        (linear2): Linear(in_features=2048, out_features=640, bias=True)\n",
      "        (norm1): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (norm2): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (norm3): LayerNorm((640,), eps=1e-05, elementwise_affine=True)\n",
      "        (dropout1): Dropout(p=0.2, inplace=False)\n",
      "        (dropout2): Dropout(p=0.2, inplace=False)\n",
      "        (dropout3): Dropout(p=0.2, inplace=False)\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (final_prediction): Sequential(\n",
      "    (0): Linear(in_features=640, out_features=256, bias=True)\n",
      "    (1): LeakyReLU(negative_slope=0.01)\n",
      "    (2): Linear(in_features=256, out_features=18, bias=True)\n",
      "    (3): Tanh()\n",
      "  )\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "dict_attr = type('dict_attr', (object,), {})\n",
    "args = dict_attr()\n",
    "args.lr = 2e-4\n",
    "args.embeddim = 640\n",
    "args.hiddendim = 256\n",
    "args.numheads = 4\n",
    "args.numlayers = 3\n",
    "args.dropout = 0.2\n",
    "args.tol = 5\n",
    "\n",
    "model = TorsionalAnglesTransformerDecoder(embed_dim=args.embeddim, hidden_dim=args.hiddendim, num_heads=args.numheads, num_layers=args.numlayers, dropout=args.dropout)\n",
    "device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "model.to(device)\n",
    "model.load_state_dict(torch.load(f\"/home2/sriram.devata/rna_project/rna_transformer/checkpoints/best_model_rna_transformer_{args.lr}_{args.embeddim}_{args.hiddendim}_{args.numheads}_{args.numlayers}_{args.dropout}_{args.tol}.pkl\", map_location=device))\n",
    "model.eval()\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ------------------- Helper functions -------------------\n",
    "\n",
    "all_dihedral_angle_names = [\"alpha\" ,\"beta\" ,\"gamma\" ,\"delta\" ,\"chi\" ,\"epsilon\" ,\"zeta\" ,\"eta\" ,\"theta\"]\n",
    "\n",
    "def convert_to_rad_and_back_to_degree(cur_res_spot_rna_pred_angles):\n",
    "\n",
    "        final_degree_angles = list()\n",
    "        for each_angle in cur_res_spot_rna_pred_angles:\n",
    "                torsion_angle = torch.tensor([each_angle * math.pi / 180])\n",
    "                final_degree_angle = torch.arctan2(torch.sin(torsion_angle),torch.cos(torsion_angle)).item() * 180 / math.pi\n",
    "                final_degree_angles.append(final_degree_angle)\n",
    "\n",
    "        return final_degree_angles\n",
    "\n",
    "def make_rad_angle_tensor_to_deg_angle_list(rad_angle_tensor):\n",
    "\n",
    "        constructed_angles = list()\n",
    "        for each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "                cos_angle = rad_angle_tensor[2*each_angle_idx]\n",
    "                sin_angle = rad_angle_tensor[2*each_angle_idx+1]\n",
    "                rad_angle = torch.arctan2(sin_angle,cos_angle)\n",
    "                deg_angle = rad_angle.item() * 180 / math.pi\n",
    "                constructed_angles.append(deg_angle)\n",
    "\n",
    "        return constructed_angles\n",
    "\n",
    "def calculate_mae(predicted, groundtruth, prediction_method=\"Unknown\"):\n",
    "\n",
    "        angle_errors_to_export = list()\n",
    "        all_angle_errors = list()\n",
    "        for each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "                all_angle_errors.append(list())\n",
    "                angle_errors_to_export.append(list())\n",
    "\n",
    "        for each_rna_idx in range(len(predicted)):\n",
    "\n",
    "                for each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "                        angle_errors_to_export[each_angle_idx].append(list())\n",
    "\n",
    "                for each_residue_idx in range(len(predicted[each_rna_idx])):\n",
    "\n",
    "                        predicted_angles = predicted[each_rna_idx][each_residue_idx]\n",
    "                        gt_angles = groundtruth[each_rna_idx][each_residue_idx]\n",
    "\n",
    "                        for each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "\n",
    "                                if prediction_method == \"Random Baseline\":\n",
    "                                        difference = abs(predicted_angles[each_angle_idx] - gt_angles[each_angle_idx])\n",
    "                                        difference = sum(difference)/len(difference)\n",
    "                                        difference = difference.item()\n",
    "                                else:\n",
    "                                        difference = abs(predicted_angles[each_angle_idx] - gt_angles[each_angle_idx])\n",
    "\n",
    "                                difference = min(difference, 360-difference)\n",
    "                                if math.isnan(difference):\n",
    "                                        continue\n",
    "                                all_angle_errors[each_angle_idx].append(difference)\n",
    "                                angle_errors_to_export[each_angle_idx][-1].append(difference)\n",
    "\n",
    "        return_maes = list()\n",
    "        all_errors = list()\n",
    "\n",
    "#         print(f\"\\n----------\\nMAEs for all dihedral angles predicted by {prediction_method}\")\n",
    "        for each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "                mae = sum(all_angle_errors[each_angle_idx])/len(all_angle_errors[each_angle_idx])\n",
    "                return_maes.append(mae)\n",
    "                all_errors.append(all_angle_errors[each_angle_idx])\n",
    "#                 print(f\"{all_dihedral_angle_names[each_angle_idx]}: {mae:.3f}\")\n",
    "\n",
    "        return return_maes, all_errors, angle_errors_to_export\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp19 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 31.50 } & \\textbf{ 12.79 } & \\textbf{ 23.72 } & 13.34  & \\textbf{ 12.71 } & \\textbf{ 11.38 } & \\textbf{ 16.47 } & 19.68  & 25.57 \\\\\n",
      "\\text{ 19\\_RW3D } & 45.63  & 19.24  & 30.84  & 13.93  & 16.98  & 19.72  & 26.17  & 28.59  & 33.63 \\\\\n",
      "\\text{ 19\\_Das\\_Human } & 39.07  & 23.32  & 30.23  & \\textbf{ 11.05 } & 12.81  & 19.99  & 25.18  & 20.19  & 29.86 \\\\\n",
      "\\text{ 19\\_SimRNA } & 50.34  & 21.96  & 33.75  & 15.13  & 19.00  & 22.39  & 44.70  & 33.98  & 40.79 \\\\\n",
      "\\text{ 19\\_Ding\\_Human } & 37.55  & 18.05  & 28.31  & 16.77  & 15.15  & 13.41  & 20.80  & 23.42  & 31.44 \\\\\n",
      "\\text{ 19\\_3dRNA } & 64.21  & 37.15  & 53.03  & 30.67  & 18.18  & 38.71  & 42.04  & 33.17  & 38.45 \\\\\n",
      "\\text{ 19\\_Dokholyan } & 49.47  & 20.29  & 51.30  & 14.85  & 18.40  & 12.67  & 23.71  & 29.31  & 34.16 \\\\\n",
      "\\text{ 19\\_LeeServer } & 34.73  & 17.87  & 28.32  & 25.31  & 24.20  & 29.19  & 24.57  & 25.22  & 33.43 \\\\\n",
      "\\text{ 19\\_Chen\\_Human } & 46.04  & 26.70  & 31.61  & 12.76  & 19.13  & 31.90  & 40.17  & 20.00  & \\textbf{ 22.34 }\\\\\n",
      "\\text{ 19\\_RNAComposer\\_Human } & 42.44  & 26.11  & 36.06  & 12.56  & 18.59  & 17.17  & 32.34  & 21.37  & 38.04 \\\\\n",
      "\\text{ 19\\_RNAComposer } & 44.57  & 27.30  & 38.11  & 13.70  & 18.07  & 16.53  & 32.79  & \\textbf{ 19.53 } & 35.78 \\\\\n",
      "\\text{ 19\\_Bujnicki\\_Human } & 39.57  & 18.03  & 36.00  & 16.88  & 15.30  & 18.23  & 20.90  & 22.64  & 27.44 \\\\\n",
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp21 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 53.03 } & \\textbf{ 22.00 } & \\textbf{ 41.33 } & 18.02  & \\textbf{ 20.91 } & \\textbf{ 18.68 } & \\textbf{ 37.12 } & \\textbf{ 19.77 } & 35.66 \\\\\n",
      "\\text{ 21\\_Das } & 61.72  & 31.70  & 53.55  & \\textbf{ 15.39 } & 20.99  & 23.91  & 41.21  & 26.72  & \\textbf{ 28.30 }\\\\\n",
      "\\text{ 21\\_Sanbonmatsu } & 61.46  & 45.81  & 56.49  & 29.20  & 29.01  & 42.09  & 45.34  & 31.58  & 55.01 \\\\\n",
      "\\text{ 21\\_RNAComposer } & 68.03  & 36.22  & 65.71  & 22.83  & 32.43  & 28.24  & 51.74  & 36.60  & 48.72 \\\\\n",
      "\\text{ 21\\_ChenHighLig } & 57.25  & 24.79  & 43.52  & 19.22  & 25.20  & 28.32  & 43.28  & 25.30  & 38.02 \\\\\n",
      "\\text{ 21\\_Bujnicki } & 65.21  & 27.26  & 46.95  & 24.02  & 25.26  & 24.38  & 42.75  & 29.74  & 40.39 \\\\\n",
      "\\text{ 21\\_RW3D } & 70.28  & 29.41  & 54.25  & 18.59  & 23.43  & 28.03  & 46.50  & 35.31  & 47.68 \\\\\n",
      "\\text{ 21\\_ChenLowLig } & 57.27  & 24.77  & 43.52  & 19.21  & 25.19  & 28.34  & 43.26  & 25.30  & 38.02 \\\\\n",
      "\\text{ 21\\_simRNA } & 58.74  & 27.14  & 55.74  & 18.88  & 22.48  & 26.36  & 46.95  & 32.88  & 44.75 \\\\\n",
      "\\text{ 21\\_DasLORES } & 61.72  & 31.70  & 53.55  & \\textbf{ 15.39 } & 20.99  & 23.91  & 41.21  & 26.72  & \\textbf{ 28.30 }\\\\\n",
      "\\text{ 21\\_Adamiak } & 74.02  & 25.64  & 58.76  & 21.00  & 25.06  & 27.79  & 50.23  & 32.42  & 43.39 \\\\\n",
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp15 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 22.04 } & \\textbf{ 10.48 } & \\textbf{ 17.74 } & 11.64  & \\textbf{ 10.71 } & \\textbf{ 11.58 } & 23.83  & 20.01  & 23.71 \\\\\n",
      "\\text{ 15\\_RNAComposerAS1 } & 25.71  & 15.97  & 19.46  & \\textbf{ 7.89 } & 12.90  & 13.94  & \\textbf{ 21.49 } & \\textbf{ 14.05 } & 19.74 \\\\\n",
      "\\text{ 15\\_Chen } & 30.79  & 15.84  & 18.88  & 9.26  & 15.82  & 20.79  & 29.55  & 16.27  & \\textbf{ 18.51 }\\\\\n",
      "\\text{ 15\\_RW3DAS1 } & 42.33  & 25.82  & 23.31  & 10.03  & 13.05  & 17.53  & 27.53  & 21.20  & 28.34 \\\\\n",
      "\\text{ 15\\_Adamiak } & 23.85  & 18.00  & 24.57  & 8.47  & 15.42  & 17.75  & 27.29  & 15.23  & 24.71 \\\\\n",
      "\\text{ 15\\_SimRNAAS1 } & 51.04  & 23.12  & 27.81  & 12.48  & 11.61  & 18.40  & 34.12  & 30.04  & 29.17 \\\\\n",
      "\\text{ 15\\_SimRNAAS2 } & 45.00  & 24.87  & 31.28  & 10.96  & 13.47  & 22.10  & 36.72  & 27.40  & 33.25 \\\\\n",
      "\\text{ 15\\_RW3DAS2 } & 36.70  & 26.80  & 28.95  & 10.57  & 13.74  & 20.31  & 28.91  & 25.14  & 31.50 \\\\\n",
      "\\text{ 15\\_RNAComposerAS2 } & 33.55  & 18.05  & 26.23  & 11.78  & 17.45  & 17.01  & 33.29  & 23.98  & 25.96 \\\\\n",
      "\\text{ 15\\_3dRNAAS2 } & 57.32  & 29.09  & 56.08  & 27.89  & 24.26  & 38.73  & 49.05  & 32.37  & 40.11 \\\\\n",
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp18 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 28.72 } & \\textbf{ 16.06 } & 25.11  & 14.42  & \\textbf{ 12.15 } & \\textbf{ 11.83 } & \\textbf{ 15.29 } & 20.75  & 22.76 \\\\\n",
      "\\text{ 18\\_3dRNA } & 58.13  & 28.02  & 41.99  & 27.54  & 20.08  & 33.63  & 29.53  & 32.87  & 29.45 \\\\\n",
      "\\text{ 18\\_Dokholyan } & 58.70  & 20.70  & 56.70  & 19.49  & 17.06  & 17.01  & 26.33  & 32.46  & 35.61 \\\\\n",
      "\\text{ 18\\_LeeASmodel } & 32.74  & 20.65  & \\textbf{ 22.23 } & 15.52  & 22.50  & 19.74  & 24.77  & 27.13  & 30.49 \\\\\n",
      "\\text{ 18\\_Chen } & 45.62  & 23.46  & 29.48  & 12.41  & 15.48  & 34.77  & 37.96  & 21.22  & 22.79 \\\\\n",
      "\\text{ 18\\_SimRNA } & 42.58  & 25.09  & 29.84  & 14.26  & 19.69  & 20.72  & 31.25  & 26.77  & 30.57 \\\\\n",
      "\\text{ 18\\_Lee } & 32.74  & 20.65  & \\textbf{ 22.23 } & 15.52  & 22.50  & 19.74  & 24.77  & 27.13  & 30.49 \\\\\n",
      "\\text{ 18\\_RW3D } & 34.46  & 21.39  & 25.49  & 12.76  & 13.45  & 13.88  & 17.69  & 22.96  & 22.27 \\\\\n",
      "\\text{ 18\\_RNAComposer } & 35.11  & 23.00  & 31.00  & 14.01  & 17.47  & 21.78  & 31.24  & 26.24  & 33.13 \\\\\n",
      "\\text{ 18\\_YagoubAli } & 59.66  & 25.32  & 41.62  & 13.64  & 17.01  & 26.69  & 33.67  & 29.17  & 37.34 \\\\\n",
      "\\text{ 18\\_Das } & 34.83  & 23.32  & 25.47  & \\textbf{ 7.50 } & 12.61  & 14.66  & 18.08  & \\textbf{ 13.65 } & \\textbf{ 12.79 }\\\\\n",
      "\\text{ 18\\_Ding } & 42.59  & 22.21  & 33.80  & 17.38  & 18.65  & 15.64  & 21.13  & 28.01  & 28.68 \\\\\n",
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp07 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 38.87 } & \\textbf{ 17.76 } & \\textbf{ 28.41 } & 11.69  & \\textbf{ 14.41 } & \\textbf{ 13.62 } & \\textbf{ 23.66 } & 21.46  & \\textbf{ 25.11 }\\\\\n",
      "\\text{ 7\\_Adamiak } & 56.06  & 31.90  & 43.78  & 15.18  & 20.46  & 20.70  & 38.97  & 27.76  & 37.74 \\\\\n",
      "\\text{ 7\\_Ding } & 49.32  & 20.82  & 44.12  & 13.80  & 21.24  & 21.57  & 36.47  & 29.91  & 36.74 \\\\\n",
      "\\text{ 7\\_Dokholyan } & 59.24  & 21.73  & 55.14  & 17.86  & 20.25  & 18.49  & 33.31  & 35.54  & 35.88 \\\\\n",
      "\\text{ 7\\_Chen } & 40.27  & 23.02  & 31.37  & \\textbf{ 9.70 } & 15.49  & 21.29  & 32.23  & \\textbf{ 18.60 } & 30.01 \\\\\n",
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp11 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 32.31 } & \\textbf{ 15.88 } & \\textbf{ 29.29 } & 8.71  & \\textbf{ 12.79 } & \\textbf{ 10.83 } & \\textbf{ 15.63 } & \\textbf{ 13.97 } & 23.56 \\\\\n",
      "\\text{ 11\\_Xiao } & 56.07  & 24.57  & 39.79  & 11.95  & 14.26  & 28.47  & 31.87  & 24.94  & 28.56 \\\\\n",
      "\\text{ 11\\_Ding } & 48.16  & 18.68  & 45.85  & 12.53  & 15.99  & 21.44  & 25.54  & 24.17  & 33.19 \\\\\n",
      "\\text{ 11\\_Adamiak } & 54.68  & 30.76  & 47.30  & \\textbf{ 7.51 } & 19.18  & 25.06  & 29.17  & 24.16  & 30.47 \\\\\n",
      "\\text{ 11\\_Das } & 36.17  & 21.53  & 32.38  & 7.73  & 16.29  & 15.45  & 23.47  & 14.45  & \\textbf{ 21.74 }\\\\\n",
      "\\text{ 11\\_Chen } & 38.62  & 21.74  & 30.69  & 12.21  & 15.61  & 26.63  & 35.87  & 28.12  & 28.98 \\\\\n",
      "\\text{ 11\\_Bujnicki } & 35.26  & 18.47  & 37.21  & 14.92  & 16.19  & 18.44  & 24.22  & 18.13  & 33.06 \\\\\n",
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp04 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & 41.14  & 20.47  & 33.41  & 6.97  & 15.23  & 23.42  & 22.84  & 19.52  & 22.05 \\\\\n",
      "\\text{ 4\\_santalucia } & 41.44  & 26.44  & 31.04  & 6.70  & 19.27  & 26.54  & 26.58  & 9.08  & 10.96 \\\\\n",
      "\\text{ 4\\_adamiak } & 45.71  & 32.24  & 36.41  & 6.65  & 19.12  & 24.92  & 26.77  & 12.33  & 17.01 \\\\\n",
      "\\text{ 4\\_major } & 55.23  & 32.11  & 36.95  & 11.92  & 17.72  & 30.68  & 26.94  & 15.30  & 15.30 \\\\\n",
      "\\text{ 4\\_das } & \\textbf{ 34.02 } & 21.59  & 25.63  & \\textbf{ 6.22 } & 13.74  & \\textbf{ 19.14 } & 20.62  & \\textbf{ 7.05 } & \\textbf{ 8.95 }\\\\\n",
      "\\text{ 4\\_mikolajczak } & 51.90  & 27.88  & 40.37  & 7.22  & 19.84  & 29.52  & 35.05  & 18.52  & 22.90 \\\\\n",
      "\\text{ 4\\_bujnicki } & 36.84  & 22.39  & 29.26  & 6.35  & 16.63  & 25.32  & 25.16  & 10.01  & 14.90 \\\\\n",
      "\\text{ 4\\_dokholyan } & 38.80  & 19.97  & 31.64  & 6.53  & 15.47  & 20.03  & 21.28  & 11.94  & 14.34 \\\\\n",
      "\\text{ 4\\_chen } & 34.59  & \\textbf{ 18.60 } & \\textbf{ 25.49 } & 6.47  & \\textbf{ 12.35 } & 20.39  & \\textbf{ 19.87 } & 7.52  & 9.52 \\\\\n",
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp09 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 34.52 } & \\textbf{ 13.74 } & \\textbf{ 29.98 } & 18.50  & 17.88  & 16.71  & 35.39  & 27.33  & 35.91 \\\\\n",
      "\\text{ 9\\_Bujnicki } & 48.74  & 23.07  & 40.30  & 24.92  & 23.29  & 26.10  & 33.81  & 33.70  & 34.71 \\\\\n",
      "\\text{ 9\\_Das } & 44.02  & 20.05  & 32.80  & \\textbf{ 11.08 } & \\textbf{ 15.85 } & \\textbf{ 16.34 } & \\textbf{ 21.27 } & \\textbf{ 18.15 } & \\textbf{ 25.49 }\\\\\n",
      "\\text{ 9\\_Dokholyan } & 62.25  & 18.21  & 61.37  & 19.72  & 20.89  & 19.22  & 37.42  & 35.69  & 48.96 \\\\\n",
      "\\text{ 9\\_Chen } & 43.37  & 18.47  & 33.57  & 14.51  & 20.37  & 18.57  & 23.62  & 28.39  & 30.12 \\\\\n",
      "\\text{ 9\\_Ding } & 51.52  & 21.42  & 44.02  & 17.74  & 19.95  & 22.00  & 37.42  & 35.51  & 37.28 \\\\\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\\midrule\n",
      "\\multicolumn{10}{c}{ rp08 } \\\\\n",
      "\\midrule\n",
      "\\text{ TorRNA } & \\textbf{ 31.66 } & \\textbf{ 12.79 } & \\textbf{ 23.25 } & \\textbf{ 6.57 } & \\textbf{ 10.73 } & \\textbf{ 14.09 } & 27.66  & 19.83  & 23.64 \\\\\n",
      "\\text{ 8\\_Bujnicki } & 45.12  & 21.01  & 35.56  & 13.05  & 17.51  & 24.51  & 29.55  & 21.26  & 26.02 \\\\\n",
      "\\text{ 8\\_Das } & 40.19  & 20.56  & 32.50  & 6.65  & 12.08  & 20.57  & \\textbf{ 25.98 } & \\textbf{ 14.41 } & \\textbf{ 21.49 }\\\\\n",
      "\\text{ 8\\_Adamiak } & 48.42  & 24.83  & 37.34  & 9.47  & 18.05  & 24.29  & 34.31  & 19.72  & 30.67 \\\\\n",
      "\\text{ 8\\_Ding } & 49.11  & 20.79  & 38.26  & 11.82  & 14.75  & 19.63  & 38.18  & 32.91  & 34.27 \\\\\n",
      "\\text{ 8\\_Chen } & 48.50  & 22.54  & 37.30  & 9.94  & 15.48  & 25.51  & 33.81  & 23.39  & 26.84 \\\\\n",
      "\\text{ 8\\_Dokholyan } & 66.87  & 22.16  & 59.96  & 10.56  & 17.25  & 19.64  & 41.87  & 29.14  & 44.06 \\\\\n"
     ]
    }
   ],
   "source": [
    "rna_puzzles_directory = \"/home2/sriram.devata/rna_project/standardized_dataset_rna_puzzles/\"\n",
    "for puzzle_submission_idx, each_dir in enumerate(os.listdir(rna_puzzles_directory)):\n",
    "    if not os.path.isdir(f\"{rna_puzzles_directory}/{each_dir}\"):\n",
    "        continue\n",
    "    \n",
    "    # print(f\"------Puzzle: {each_dir}------\")\n",
    "    \n",
    "    cur_puzzle_directory = f\"{rna_puzzles_directory}/{each_dir}\"\n",
    "    \n",
    "    all_puzzle_submissions = os.listdir(cur_puzzle_directory)\n",
    "    puzzle_solution_file = None\n",
    "    for each_puzzle_submission in all_puzzle_submissions:\n",
    "        if \"solution\" in each_puzzle_submission and \"pdb\" in each_puzzle_submission:\n",
    "            puzzle_solution_file = each_puzzle_submission\n",
    "    if puzzle_solution_file is None:\n",
    "        continue\n",
    "    # print(puzzle_solution_file)\n",
    "        \n",
    "    # get DSSR angles\n",
    "    puzzle_solution_rna_object = RNA(puzzle_solution_file, calc_rna_fm_embeddings=True, load_dssr_dihedrals=True, load_coords=False, dssr_path=cur_puzzle_directory)\n",
    "    gt_torsional_angles = puzzle_solution_rna_object.dssr_torsion_angles.unsqueeze(0).to(device)\n",
    "    all_gt_angles = list()\n",
    "    for each_residue_idx in range(len(gt_torsional_angles[0])):\n",
    "        cur_res_gt_angles = make_rad_angle_tensor_to_deg_angle_list(gt_torsional_angles[0][each_residue_idx])\n",
    "        all_gt_angles.append(cur_res_gt_angles)\n",
    "        \n",
    "    # --- Predict the dihedral angles with TorRNA ---\n",
    "    rna_fm_embeddings = puzzle_solution_rna_object.rna_fm_embeddings.unsqueeze(0).to(device)\n",
    "    initial_embeddings = puzzle_solution_rna_object.rna_fm_initial_embeddings.unsqueeze(0).to(device)\n",
    "    gt_torsional_angles = puzzle_solution_rna_object.dssr_torsion_angles.unsqueeze(0).to(device)\n",
    "    padding_mask = torch.sum(torch.abs(gt_torsional_angles), dim=-1) < 1e-5\n",
    "    output = model(rna_fm_embeddings, padding_mask, initial_embeddings).detach().cpu()[0]\n",
    "    all_predicted_angles = list()\n",
    "    for each_residue_idx in range(len(output)):\n",
    "        if padding_mask[0][each_residue_idx]:\n",
    "            continue\t# True if this residue is padded\n",
    "        predicted_angles = make_rad_angle_tensor_to_deg_angle_list(output[each_residue_idx])\n",
    "        all_predicted_angles.append(predicted_angles)\n",
    "    all_maes, all_predicted_errors, angle_errors_to_export = calculate_mae([all_predicted_angles], [all_gt_angles], prediction_method=\"TorRNA\")\n",
    "    # print(f\"TorRNA MAEs: {all_maes}\")\n",
    "\n",
    "\n",
    "    all_submissions_rmsds = pd.read_csv(f\"{rna_puzzles_directory}/{each_dir}/rmsds.csv\")\n",
    "    all_groups = list(all_submissions_rmsds[\"fn\"])\n",
    "    all_groups = ['_'.join(x.split('_')[:-2]) for x in all_groups if \"solution\" not in x]\n",
    "    all_groups = list(set(all_groups))\n",
    "    \n",
    "    each_puzzle_all_groups_labels = [\"TorRNA\"]\n",
    "    each_puzzle_all_groups_maes = [all_maes]\n",
    "    # get the RMSDs of all submissions by this group\n",
    "    for each_group in all_groups:\n",
    "        each_group_all_rmsds = all_submissions_rmsds\n",
    "        each_group_all_rmsds = all_submissions_rmsds.loc[all_submissions_rmsds['fn'].str.contains(each_group)]\n",
    "        \n",
    "        # get the filename of the submission with the lowest RMSD\n",
    "        each_group_best_pdb = each_group_all_rmsds.nsmallest(n=1, columns=\"rmsd_all\").iloc[0][\"fn\"]\n",
    "        each_group_best_rna_object = RNA(each_group_best_pdb, calc_rna_fm_embeddings=True, load_dssr_dihedrals=True, load_coords=False, dssr_path=cur_puzzle_directory)\n",
    "        \n",
    "        # if there is a chain mismatch, no point comparing the predictions\n",
    "        if puzzle_solution_rna_object.dssr_full_seq != each_group_best_rna_object.dssr_full_seq:\n",
    "            continue\n",
    "            \n",
    "        # get the DSSR angles of each of the group submission\n",
    "        each_group_best_angles = each_group_best_rna_object.dssr_torsion_angles.unsqueeze(0).to(device)\n",
    "        all_each_group_best_angles = list()\n",
    "        for each_residue_idx in range(len(each_group_best_angles[0])):\n",
    "            cur_res_gt_angles = make_rad_angle_tensor_to_deg_angle_list(each_group_best_angles[0][each_residue_idx])\n",
    "            all_each_group_best_angles.append(cur_res_gt_angles)\n",
    "        all_maes, all_predicted_errors, angle_errors_to_export = calculate_mae([all_each_group_best_angles], [all_gt_angles], prediction_method=\"TorRNA\")\n",
    "        # print(f\"{each_group}: {all_maes}\")\n",
    "        if sum(all_maes) < 0.001:\n",
    "            continue # Some cases have all 0 MAEs, it's an error\n",
    "        each_puzzle_all_groups_labels.append(each_group)\n",
    "        each_puzzle_all_groups_maes.append(all_maes)\n",
    "\n",
    "        \n",
    "    if len(each_puzzle_all_groups_maes) <= 3:\n",
    "        continue\n",
    "    # print(f\"------Puzzle: {each_dir}------\")\n",
    "    print(f\"\\midrule\")\n",
    "    # print(f\"{each_dir} & & & & & & & & & \\\\\\\\\")\n",
    "    print(f\"\\\\multicolumn{{10}}{{c}}{{ {each_dir} }} \\\\\\\\\")\n",
    "    print(f\"\\midrule\")\n",
    "    for each_group_name, each_group_mae in zip(each_puzzle_all_groups_labels,each_puzzle_all_groups_maes):\n",
    "        each_group_name = each_group_name.replace('_', '\\_')\n",
    "        print(f\"\\\\text{{ {each_group_name} }}\", end='')\n",
    "        for mae_idx, each_mae in enumerate(each_group_mae):\n",
    "            \n",
    "            if each_mae == min([x[mae_idx] for x in each_puzzle_all_groups_maes]):\n",
    "                print(f\" & \\\\textbf{{ {each_mae:.2f} }}\", end='')\n",
    "            else:\n",
    "                print(f\" & {each_mae:.2f} \", end='')\n",
    "            \n",
    "        print(\"\\\\\\\\\")\n",
    "    \n",
    "    \n",
    "    if puzzle_submission_idx > 30:\n",
    "        break\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "torrna",
   "language": "python",
   "name": "torrna"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}