Predict Torsional Angles for Residues in Different Regions.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import math\n",
    "from tqdm import tqdm\n",
    "import wandb\n",
    "import pickle\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": {},
   "outputs": [],
   "source": [
    "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",
    "\tfinal_degree_angles = list()\n",
    "\tfor each_angle in cur_res_spot_rna_pred_angles:\n",
    "\t\ttorsion_angle = torch.tensor([each_angle * math.pi / 180])\n",
    "\t\tfinal_degree_angle = torch.arctan2(torch.sin(torsion_angle),torch.cos(torsion_angle)).item() * 180 / math.pi\n",
    "\t\tfinal_degree_angles.append(final_degree_angle)\n",
    "\n",
    "\treturn final_degree_angles\n",
    "\n",
    "def make_rad_angle_tensor_to_deg_angle_list(rad_angle_tensor):\n",
    "\n",
    "\tconstructed_angles = list()\n",
    "\tfor each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "\t\tcos_angle = rad_angle_tensor[2*each_angle_idx]\n",
    "\t\tsin_angle = rad_angle_tensor[2*each_angle_idx+1]\n",
    "\t\trad_angle = torch.arctan2(sin_angle,cos_angle)\n",
    "\t\tdeg_angle = rad_angle.item() * 180 / math.pi\n",
    "\t\tconstructed_angles.append(deg_angle)\n",
    "\n",
    "\treturn constructed_angles\n",
    "\n",
    "def calculate_mae(predicted, groundtruth, prediction_method=\"Unknown\"):\n",
    "\n",
    "\tall_angle_errors = list()\n",
    "\tfor each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "\t\tall_angle_errors.append(list())\n",
    "\n",
    "\tfor each_rna_idx in range(len(predicted)):\n",
    "\t\tfor each_residue_idx in range(len(predicted[each_rna_idx])):\n",
    "\n",
    "\t\t\tpredicted_angles = predicted[each_rna_idx][each_residue_idx]\n",
    "\t\t\tgt_angles = groundtruth[each_rna_idx][each_residue_idx]\n",
    "\n",
    "\t\t\tfor each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "\n",
    "\t\t\t\tif prediction_method == \"Random Baseline\":\n",
    "\t\t\t\t\tdifference = abs(predicted_angles[each_angle_idx] - gt_angles[each_angle_idx])\n",
    "\t\t\t\t\tdifference = sum(difference)/len(difference)\n",
    "\t\t\t\t\tdifference = difference.item()\n",
    "\t\t\t\telse:\n",
    "\t\t\t\t\tdifference = abs(predicted_angles[each_angle_idx] - gt_angles[each_angle_idx])\n",
    "\n",
    "\t\t\t\tdifference = min(difference, 360-difference)\n",
    "\t\t\t\tif math.isnan(difference):\n",
    "\t\t\t\t\tcontinue\n",
    "\t\t\t\tall_angle_errors[each_angle_idx].append(difference)\n",
    "\n",
    "\treturn_maes = list()\n",
    "\tall_errors = list()\n",
    "\n",
    "\tprint(f\"\\n----------\\nMAEs for all dihedral angles predicted by {prediction_method}\")\n",
    "\tfor each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "\t\tmae = sum(all_angle_errors[each_angle_idx])/len(all_angle_errors[each_angle_idx])\n",
    "\t\treturn_maes.append(mae)\n",
    "\t\tall_errors.append(all_angle_errors[each_angle_idx])\n",
    "\t\tprint(f\"{all_dihedral_angle_names[each_angle_idx]}: {mae:.3f}\")\n",
    "\n",
    "\treturn return_maes, all_errors\n",
    "\n",
    "def calculate_mae_for_positions(predicted, groundtruth, nuc_positions, prediction_method=\"\"):\n",
    "    all_angle_errors = list()\n",
    "    for each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "        all_angle_errors.append(list())\n",
    "\n",
    "    for each_residue_idx in nuc_positions:\n",
    "        \n",
    "        predicted_angles = predicted[each_residue_idx]\n",
    "        gt_angles = groundtruth[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",
    "            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",
    "\n",
    "    return_maes = list()\n",
    "    all_errors = list()\n",
    "\n",
    "    for each_angle_idx in range(len(all_dihedral_angle_names)):\n",
    "        try:\n",
    "            mae = sum(all_angle_errors[each_angle_idx])/len(all_angle_errors[each_angle_idx])\n",
    "        except:\n",
    "            mae = 0\n",
    "        return_maes.append(mae)\n",
    "        all_errors.append(all_angle_errors[each_angle_idx])\n",
    "\n",
    "    return return_maes, all_errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "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",
       ")"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "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",
    "# Setting up the transformer model\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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdbs_path=\"/home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs/\"\n",
    "dssr_path=pdbs_path\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",
    "\ttraining_pdbs, validation_pdbs, testing_pdbs = pickle.load(fp)\n",
    "training_files = list()\n",
    "for each_training_pdb in training_pdbs:\n",
    "\ttraining_files.append(f\"{pdbs_path}/{each_training_pdb}.pdb\")\n",
    "testing_files = list()\n",
    "for each_testing_pdb in testing_pdbs:\n",
    "\ttesting_files.append(f\"{pdbs_path}/{each_testing_pdb}.pdb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def get_all_regions(each_test_pdb, pdbs_path):\n",
    "    each_pdb_file = f\"{pdbs_path}/{each_test_pdb}.pdb\"\n",
    "    try:\n",
    "        rna_object = RNA(each_pdb_file, calc_rna_fm_embeddings=True, load_dssr_dihedrals=True, load_coords=False)\n",
    "    except:\n",
    "        return None\n",
    "    # print(pdb_idx, rna_object.dssr_full_seq, len(rna_object.dssr_full_seq))\n",
    "    \n",
    "    \n",
    "    # --- Predict the dihedral angles ---\n",
    "    rna_fm_embeddings = rna_object.rna_fm_embeddings.unsqueeze(0).to(device)\n",
    "    initial_embeddings = rna_object.rna_fm_initial_embeddings.unsqueeze(0).to(device)\n",
    "    pading_mask_shape = list(rna_fm_embeddings.shape[:-1])\n",
    "    pading_mask_shape[-1] -= 2\n",
    "    padding_mask = torch.zeros(pading_mask_shape, device=device, dtype=torch.bool)\n",
    "\n",
    "    # Get the dihedral angle predictions. This will be the \"ground truth\" dihedral angles\n",
    "    output = model(rna_fm_embeddings, padding_mask, initial_embeddings).detach().cpu()[0]\n",
    "    rad_gt_angles = rna_object.dssr_torsion_angles\n",
    "    all_predicted_angles = list()\n",
    "    all_gt_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",
    "        gt_angles = make_rad_angle_tensor_to_deg_angle_list(rad_gt_angles[each_residue_idx])\n",
    "        \n",
    "        all_predicted_angles.append(predicted_angles)\n",
    "        all_gt_angles.append(gt_angles)\n",
    "    \n",
    "    \n",
    "    each_pdb_dssr = f\"{dssr_path}/{each_test_pdb}.out\"\n",
    "    with open(each_pdb_dssr, 'r') as f:\n",
    "        each_pdb_dssr_lines = f.readlines()\n",
    "        each_pdb_dssr_lines = [x.strip() for x in each_pdb_dssr_lines]\n",
    "        \n",
    "    base_pair_lines = list()\n",
    "    for line_idx, each_line in enumerate(each_pdb_dssr_lines):\n",
    "        if \"List of\" in each_line and \"base pair\" in each_line:\n",
    "            for each_base_pair_line_idx, each_base_pair_line in enumerate(each_pdb_dssr_lines[line_idx+2:]):\n",
    "                if \"****\" in each_base_pair_line:\n",
    "                    break\n",
    "                if len(each_base_pair_line) > 0:\n",
    "                    base_pair_lines.append(each_base_pair_line)\n",
    "            break\n",
    "    \n",
    "    all_base_pairs = list() # Has all base pairs. Is of the form `[(1,15), (2,14), ...]`\n",
    "    for each_base_pair_line in base_pair_lines:\n",
    "        # each_base_pair_line is of the form: `1 A.G1   A.C15    G-C WC    19-XIX    cWW  cW-W`\n",
    "        each_base_pair_line = each_base_pair_line.split()\n",
    "        bp_first = each_base_pair_line[1]\n",
    "        bp_second = each_base_pair_line[2]\n",
    "        \n",
    "        # Get the base pair type as well\n",
    "        bp_type = each_base_pair_line[4]\n",
    "        \n",
    "        bp_first_num = int(''.join(filter(str.isdigit, bp_first)))  # Remove all non-digits from each string of the form `A.G1`\n",
    "        bp_second_num = int(''.join(filter(str.isdigit, bp_second)))\n",
    "        \n",
    "        all_base_pairs.append((bp_first_num, bp_second_num, bp_type))\n",
    "        \n",
    "            \n",
    "    # Find the nucleotides that are non-canonical pairs\n",
    "    non_canonical_nums = list()\n",
    "    if len(all_base_pairs) > 0:\n",
    "        for each_bp in all_base_pairs:\n",
    "            bp_one, bp_two, bp_type = each_bp\n",
    "            if bp_type != \"WC\" and bp_type != \"rWC\" and bp_type != \"Wobbles\":\n",
    "                non_canonical_nums += [bp_one, bp_two]\n",
    "                \n",
    "    # Find the nucleotides that are canonical pairs\n",
    "    all_canonical_nums = list()\n",
    "    if len(all_base_pairs) > 0:\n",
    "        for each_bp in all_base_pairs:\n",
    "            bp_one, bp_two, bp_type = each_bp\n",
    "            if bp_type == \"WC\" or bp_type == \"rWC\" or bp_type == \"Wobbles\":\n",
    "                all_canonical_nums += [bp_one, bp_two]\n",
    "                \n",
    "                \n",
    "    # Find all hairpin loop nucleotides\n",
    "    hairpin_loop_lines = list() # \n",
    "    for line_idx, each_line in enumerate(each_pdb_dssr_lines):\n",
    "        if \"List of\" in each_line and \"hairpin loop\" in each_line:\n",
    "            for each_hairpin_loop_line_idx, each_hairpin_loop_line in enumerate(each_pdb_dssr_lines[line_idx+1:]):\n",
    "                if \"****\" in each_hairpin_loop_line:\n",
    "                    break\n",
    "                if len(each_hairpin_loop_line) > 0:\n",
    "                    hairpin_loop_lines.append(each_hairpin_loop_line)\n",
    "            break\n",
    "            \n",
    "    hairpin_loop_nums = list()\n",
    "    for each_hairpin_loop_line_idx, each_hairpin_loop_line in enumerate(hairpin_loop_lines):\n",
    "        if \"hairpin loop\" in each_hairpin_loop_line:\n",
    "            # cur_hairpin_loop_line is of the form `nts=7 CUCAACU A.C6,A.U7,A.C8,A.A9,A.A10,A.C11,A.U12`\n",
    "            cur_hairpin_loop_line = hairpin_loop_lines[each_hairpin_loop_line_idx+3]\n",
    "            _, _, cur_hairpin_nucs = cur_hairpin_loop_line.strip().split()\n",
    "            cur_hairpin_nucs = cur_hairpin_nucs.split(',')\n",
    "            cur_hairpin_nucs = [int(''.join(filter(str.isdigit, x))) for x in cur_hairpin_nucs]\n",
    "            # Get only the middle 4-5 nucleotides of the hairpin - the nucleotides at the very edge\n",
    "            while len(cur_hairpin_nucs) > 5:\n",
    "                cur_hairpin_nucs = cur_hairpin_nucs[1:-1]\n",
    "            \n",
    "            hairpin_loop_nums += cur_hairpin_nucs\n",
    "    \n",
    "    \n",
    "    # Find all multiplet nucleotides\n",
    "    multiplet_lines = list() # \n",
    "    for line_idx, each_line in enumerate(each_pdb_dssr_lines):\n",
    "        if \"List of\" in each_line and \"multiplet\" in each_line:\n",
    "            for each_multiplet_line_idx, each_multiplet_line in enumerate(each_pdb_dssr_lines[line_idx+1:]):\n",
    "                if \"****\" in each_multiplet_line:\n",
    "                    break\n",
    "                if len(each_multiplet_line) > 0:\n",
    "                    multiplet_lines.append(each_multiplet_line)\n",
    "            break\n",
    "\n",
    "    multiplet_nums = list() # is of the form `['1 nts=3 AAU A.A4,A.A6,A.U20', '2 nts=3 UCG A.U5,A.C7,A.G19', '3 nts=3 UAA A.U8,A.A17,A.A18']`\n",
    "    for each_multiplet_line in multiplet_lines:\n",
    "        _, _, _, multiplet_nucleotides = each_multiplet_line.strip().split()\n",
    "        # multiplet_nucleotides is of the form `A.A4,A.A6,A.U20`\n",
    "        for each_multiplet_nucleotide in multiplet_nucleotides.split(','):\n",
    "            # each_multiplet_nucleotide is of the form `A.A4`. Have to remove all alphabets and special characters\n",
    "            each_multiplet_nucleotide_num = int(''.join(filter(str.isdigit, each_multiplet_nucleotide)))\n",
    "            multiplet_nums.append(each_multiplet_nucleotide_num)\n",
    "#     if len(lone_paired_nums) == 0:\n",
    "#         continue\n",
    "\n",
    "    # Get all the lines from the summary of all nucleotides\n",
    "    summary_lines = list() # \n",
    "    for line_idx, each_line in enumerate(each_pdb_dssr_lines):\n",
    "        # The summary of each nucleotide has a small note that ends with this line\n",
    "        if \"break: no backbone linkage between\" in each_line:\n",
    "            for each_summary_line_idx, each_summary_line in enumerate(each_pdb_dssr_lines[line_idx+1:]):\n",
    "                if \"****\" in each_summary_line:\n",
    "                    break\n",
    "                if len(each_summary_line) > 0:\n",
    "                    summary_lines.append(each_summary_line)\n",
    "            break\n",
    "    \n",
    "    # Find the nucleotides that are unpaired\n",
    "    unpaired_nums = list()\n",
    "    if len(all_base_pairs) > 0:\n",
    "        bp_ones, bp_twos, bp_types = zip(*all_base_pairs)\n",
    "        all_paired_nums = list(set(bp_ones + bp_twos))\n",
    "        \n",
    "        for each_summary_line in summary_lines:\n",
    "            nuc_num, nuc_letter, _, nuc_num_in_chain, _, _ = each_summary_line.strip().split()\n",
    "            nuc_num_in_chain = int(''.join(filter(str.isdigit, nuc_num_in_chain)))\n",
    "            if nuc_num_in_chain not in all_paired_nums:\n",
    "                unpaired_nums.append(nuc_num_in_chain)\n",
    "    else:\n",
    "        for each_summary_line in summary_lines:\n",
    "            nuc_num, nuc_letter, _, nuc_num_in_chain, _, _ = each_summary_line.strip().split()\n",
    "            nuc_num_in_chain = int(''.join(filter(str.isdigit, nuc_num_in_chain)))\n",
    "            unpaired_nums.append(nuc_num_in_chain)\n",
    "    \n",
    "    # Find all the pseudoknotted nucleotides\n",
    "    pseudoknot_nums = list()\n",
    "    for each_pseudoknot_line in summary_lines:\n",
    "        # Each line is of the form `1  G ( A.G1  0.008  anti,~C3'-endo,BI,canonical,non-pair-contact,helix-end,stem-end`\n",
    "        each_pseudoknot_line = each_pseudoknot_line.strip().split()\n",
    "        if each_pseudoknot_line[2] in ['[', ']', '{', '}', '<', '>']:\n",
    "            pseudoknot_nums.append(int(each_pseudoknot_line[0]))\n",
    "     \n",
    "    \n",
    "    # Find all canonical nested nucleotides\n",
    "    canonical_nested_nums = list()\n",
    "    if len(all_base_pairs) > 0:\n",
    "        bp_ones, bp_twos, bp_types = zip(*all_base_pairs)\n",
    "        all_paired_nums = list(set(bp_ones + bp_twos))\n",
    "        for each_canonical_nested_pair_line in summary_lines:\n",
    "            # Each line is of the form `1  G ( A.G1  0.008  anti,~C3'-endo,BI,canonical,non-pair-contact,helix-end,stem-end`\n",
    "            nuc_num, nuc_letter, _, _, _, each_nucleotide_attributes = each_canonical_nested_pair_line.strip().split()\n",
    "            each_nucleotide_attributes = each_nucleotide_attributes.split(',')\n",
    "            is_pseudoknot = False # Checks if this nucleotide is part of a pseudoknot\n",
    "            is_nested = False # Checks if this nucleotide is part of a nested base pair\n",
    "            for each_attribute in each_nucleotide_attributes:\n",
    "                if \"helix\" in each_attribute:\n",
    "                    is_nested = True\n",
    "                elif \"pseudoknot\" in each_attribute:\n",
    "                    is_pseudoknot = True\n",
    "            if int(nuc_num) in all_paired_nums and is_nested and not is_pseudoknot:\n",
    "                canonical_nested_nums.append(int(nuc_num))\n",
    "            \n",
    "    \n",
    "    # Find all nucleotides belonging to lone base pairs\n",
    "    # Defined as \"single tertiary base pair not associated with a helix\" from https://www.biorxiv.org/content/10.1101/677310v1.full\n",
    "    lone_base_pairs_nums = list()\n",
    "    if len(all_base_pairs) > 0:\n",
    "        bp_ones, bp_twos, bp_types = zip(*all_base_pairs)\n",
    "        all_paired_nums = list(set(bp_ones + bp_twos))\n",
    "        for each_canonical_nested_pair_line in summary_lines:\n",
    "            # Each line is of the form `1  G ( A.G1  0.008  anti,~C3'-endo,BI,canonical,non-pair-contact,helix-end,stem-end`\n",
    "            nuc_num, nuc_letter, _, _, _, each_nucleotide_attributes = each_canonical_nested_pair_line.strip().split()\n",
    "            each_nucleotide_attributes = each_nucleotide_attributes.split(',')\n",
    "            is_in_helix = False # Checks if this nucleotide is part of a helix\n",
    "            for each_attribute in each_nucleotide_attributes:\n",
    "                if \"helix\" in each_attribute or \"stem\" in each_attribute:\n",
    "                    is_in_helix = True\n",
    "            if int(nuc_num) in all_paired_nums and not is_in_helix:\n",
    "                lone_base_pairs_nums.append(int(nuc_num))        \n",
    "                \n",
    "    # The numbers until now are the nucleotide numbers 1..N, but might be some M..N if this is some chain other than A\n",
    "    # Have to change these numbers to be 0..N specific to this chain\n",
    "    nuc_num_in_chain_to_zero_index_mapping = dict()\n",
    "    for each_canonical_nested_pair_line in summary_lines:\n",
    "        # Each line is of the form `1  G ( A.G1  0.008  anti,~C3'-endo,BI,canonical,non-pair-contact,helix-end,stem-end`\n",
    "        nuc_num, nuc_letter, _, nuc_num_in_chain, _, _ = each_canonical_nested_pair_line.strip().split()\n",
    "        nuc_num = int(nuc_num) - 1 # This is originally 1..N\n",
    "        nuc_num_in_chain = int(''.join(filter(str.isdigit, nuc_num_in_chain))) # This is either 1..N or M..N depending on the chain\n",
    "        nuc_num_in_chain_to_zero_index_mapping[str(nuc_num_in_chain)] = nuc_num\n",
    "        \n",
    "    def correct_mapping_in_nuc_num(given_input_list):\n",
    "        return_list = list()\n",
    "        for each_nuc_num in given_input_list:\n",
    "            return_list.append(nuc_num_in_chain_to_zero_index_mapping[str(each_nuc_num)])\n",
    "        return return_list\n",
    "    \n",
    "\n",
    "    \n",
    "    if len(rna_object.dssr_full_seq) != len(nuc_num_in_chain_to_zero_index_mapping):\n",
    "        print(f\"SKIPPING {each_pdb_dssr} due to mismatch of number of residues\")\n",
    "\n",
    "#         print(each_pdb_dssr)\n",
    "#         print(nuc_num_in_chain_to_zero_index_mapping)\n",
    "#         print(pdb_idx, rna_object.dssr_full_seq, len(rna_object.dssr_full_seq))\n",
    "#         print(unpaired_nums)\n",
    "    \n",
    "        return None\n",
    "    \n",
    "    try:\n",
    "        unpaired_nums = correct_mapping_in_nuc_num(unpaired_nums)\n",
    "        lone_base_pairs_nums = correct_mapping_in_nuc_num(lone_base_pairs_nums)\n",
    "        pseudoknot_nums = correct_mapping_in_nuc_num(pseudoknot_nums)\n",
    "        multiplet_nums = correct_mapping_in_nuc_num(multiplet_nums)\n",
    "        non_canonical_nums = correct_mapping_in_nuc_num(non_canonical_nums)\n",
    "        all_canonical_nums = correct_mapping_in_nuc_num(all_canonical_nums)\n",
    "        hairpin_loop_nums = correct_mapping_in_nuc_num(hairpin_loop_nums)\n",
    "        canonical_nested_nums = correct_mapping_in_nuc_num(canonical_nested_nums)\n",
    "        \n",
    "        bp_ones, bp_twos, bp_types = zip(*all_base_pairs)\n",
    "        all_paired_nums = list(set(bp_ones + bp_twos))\n",
    "        \n",
    "        reshaped_all_base_pairs = list()\n",
    "        for i in range(len(bp_ones)):\n",
    "            reshaped_all_base_pairs.append(bp_ones[i])\n",
    "            reshaped_all_base_pairs.append(bp_twos[i])\n",
    "        reshaped_all_base_pairs = correct_mapping_in_nuc_num(reshaped_all_base_pairs)\n",
    "        # reshaped_all_base_pairs has all the base pairs equivalent to doing .reshape(-1) on the numbers\n",
    "        \n",
    "    except Exception as e:\n",
    "        print(f\"Unable to correct mappings for {each_pdb_dssr}\", e)\n",
    "        return None\n",
    "        \n",
    "#     print(each_pdb_file, each_pdb_dssr, len(rna_object.dssr_full_seq), rna_object.dssr_full_seq)\n",
    "#     print(f\"All base pairs: {all_base_pairs}\")\n",
    "#     print(f\"All unpaired bases: {unpaired_nums}\")\n",
    "#     print(f\"All non-canonicals: {non_canonical_nums}\")\n",
    "#     print(f\"All canonicals: {all_canonical_nums}\")\n",
    "#     print(f\"All hairpins: {hairpin_loop_nums}\")\n",
    "#     print(f\"All multiplets: {multiplet_nums}\")\n",
    "#     print(f\"All pseudoknots: {pseudoknot_nums}\")\n",
    "#     print(f\"All canonical nested: {canonical_nested_nums}\")\n",
    "#     print(f\"All lone base pairs: {lone_base_pairs_nums}\")\n",
    "\n",
    "\n",
    "    return [len(rna_object.dssr_full_seq), all_gt_angles, all_predicted_angles, nuc_num_in_chain_to_zero_index_mapping, unpaired_nums, lone_base_pairs_nums, pseudoknot_nums, multiplet_nums, non_canonical_nums, all_canonical_nums, hairpin_loop_nums, canonical_nested_nums, reshaped_all_base_pairs] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 4/172 [00:23<16:09,  5.77s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5LYU_B.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  6%|▌         | 10/172 [00:47<10:34,  3.92s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6O7H_G.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 12/172 [00:56<12:01,  4.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5WTI_B.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 18/172 [01:19<09:46,  3.81s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7UZX_G.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 20/172 [01:26<09:07,  3.60s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6BK8_E.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 12%|█▏        | 21/172 [01:30<08:54,  3.54s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//2XDB_G.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 13%|█▎        | 22/172 [01:35<10:17,  4.11s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5ZWO_G.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 14%|█▍        | 24/172 [01:42<09:14,  3.75s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1U6B_C.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 27/172 [01:54<10:04,  4.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//3WQY_C.out '18'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 35/172 [02:28<09:31,  4.17s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1C9S_W.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|██▍       | 41/172 [02:55<09:17,  4.26s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8I45_A.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 46/172 [03:12<07:28,  3.56s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5MWI_A.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▉       | 50/172 [03:30<09:10,  4.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//2DU3_D.out '18'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███▏      | 54/172 [03:46<08:27,  4.30s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5AH5_C.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 34%|███▍      | 59/172 [04:03<06:29,  3.45s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7XPL_G.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 42%|████▏     | 72/172 [04:46<04:56,  2.96s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6NUE_H.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 42%|████▏     | 73/172 [04:48<04:32,  2.76s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//2ZUE_B.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|████▍     | 76/172 [04:57<04:24,  2.76s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7WWV_M.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 81/172 [05:12<04:46,  3.15s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1F7U_B.out '18'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 49%|████▉     | 84/172 [05:22<04:46,  3.26s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//4RDX_C.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▍    | 94/172 [05:55<03:40,  2.83s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8FZT_A.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 58%|█████▊    | 100/172 [06:09<03:00,  2.51s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5EEU_W.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 60%|██████    | 104/172 [06:20<03:06,  2.74s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8AW3_1.out '55'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 66%|██████▌   | 113/172 [06:46<02:41,  2.73s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8BVM_U.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 68%|██████▊   | 117/172 [06:57<02:30,  2.74s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5OIK_P.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 69%|██████▉   | 119/172 [07:02<02:09,  2.44s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7UOB_T.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 121/172 [07:10<02:43,  3.21s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5U30_B.out '75'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 72%|███████▏  | 123/172 [07:13<02:03,  2.52s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7SBA_Z.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 73%|███████▎  | 125/172 [07:18<01:56,  2.48s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7ZHG_5.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████▌  | 129/172 [07:32<02:20,  3.27s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7K16_P.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 83%|████████▎ | 143/172 [08:11<01:28,  3.04s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1MNX_A.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 84%|████████▎ | 144/172 [08:14<01:19,  2.84s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1GTN_W.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 85%|████████▍ | 146/172 [08:19<01:08,  2.65s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5LWJ_A.out '35'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|████████▉ | 154/172 [08:43<00:51,  2.87s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5Z71_A.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 91%|█████████ | 156/172 [08:47<00:42,  2.65s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5ED1_B.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|█████████▏| 159/172 [08:57<00:40,  3.15s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1E8O_E.out '15'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 94%|█████████▍| 162/172 [09:07<00:31,  3.18s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6BSI_R.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 172/172 [09:38<00:00,  3.37s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAEs for various regions of TorRNA\n",
      "alpha,beta,gamma,delta,chi,epsilon,zeta,eta,theta,\n",
      "Unpaired (27.77\\%),61.56,29.32,48.21,22.45,32.14,26.75,46.76,48.18,48.83,\n",
      "Lone Pairs (5.99\\%),44.43,22.18,37.18,16.28,22.48,22.96,42.93,25.10,36.16,\n",
      "Pseudoknots (2.71\\%),38.30,18.51,30.23,11.88,13.38,18.73,27.96,23.60,27.61,\n",
      "Multiplets (9.24\\%),50.89,24.89,42.53,17.36,25.07,22.88,42.60,27.60,37.04,\n",
      "Non-Canonical Pairs (25.61\\%),47.24,24.14,41.32,15.58,24.48,22.42,37.94,25.94,33.89,\n",
      "All Canonical Pairs (50.51\\%),29.51,15.06,27.54,7.82,9.37,13.80,15.18,11.10,14.43,\n",
      "Canonical Nested Pairs (46.49\\%),33.32,16.89,30.45,9.01,12.43,15.13,18.51,13.95,17.30,\n",
      "Hairpin Loops (11.72\\%),62.57,26.55,42.77,22.75,28.52,26.57,46.18,48.98,45.43,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "total_num_nucleotides = 0\n",
    "\n",
    "unpaired_errors = [list() for _ in range(9)]\n",
    "lone_pairs_errors = [list() for _ in range(9)]\n",
    "pseudoknots_errors = [list() for _ in range(9)]\n",
    "multiplets_errors = [list() for _ in range(9)]\n",
    "noncanonical_pairs_errors = [list() for _ in range(9)]\n",
    "all_canonical_pairs_errors = [list() for _ in range(9)]\n",
    "canonical_nested_pairs_errors = [list() for _ in range(9)]\n",
    "hairpin_loop_errors = [list() for _ in range(9)]\n",
    "\n",
    "for pdb_idx, each_test_pdb in tqdm(enumerate(testing_pdbs), total=len(testing_pdbs)):\n",
    "    \n",
    "    regions_output = get_all_regions(each_test_pdb, pdbs_path)\n",
    "    if regions_output is not None:\n",
    "        len_seq, all_gt_angles, all_predicted_angles, nuc_num_in_chain_to_zero_index_mapping, unpaired_nums, lone_base_pairs_nums, pseudoknot_nums, multiplet_nums, non_canonical_nums, all_canonical_nums, hairpin_loop_nums, canonical_nested_nums, reshaped_all_base_pairs = regions_output\n",
    "    else:\n",
    "        continue\n",
    "        \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, unpaired_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        unpaired_errors[angle_idx] += all_errors[angle_idx]\n",
    "    \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, lone_base_pairs_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        lone_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "    \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, pseudoknot_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        pseudoknots_errors[angle_idx] += all_errors[angle_idx]\n",
    "    \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, multiplet_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        multiplets_errors[angle_idx] += all_errors[angle_idx]\n",
    "    \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, non_canonical_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        noncanonical_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "        \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, all_canonical_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        all_canonical_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "    \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, canonical_nested_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        canonical_nested_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "        \n",
    "    return_maes, all_errors = calculate_mae_for_positions(all_predicted_angles, all_gt_angles, hairpin_loop_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "    for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "        hairpin_loop_errors[angle_idx] += all_errors[angle_idx]\n",
    "        \n",
    "    total_num_nucleotides += len_seq\n",
    "        \n",
    "    if pdb_idx > 1000:\n",
    "        break\n",
    "\n",
    "def print_region_maes(region_errors, region_name=\"Unknown\"):\n",
    "    print(f\"{region_name} ({len(region_errors[0])/total_num_nucleotides*100:.2f}\\%)\", end=',')\n",
    "    for each_angle_idx, each_angle_errors in enumerate(region_errors):\n",
    "        try:\n",
    "            mae = sum(each_angle_errors)/len(each_angle_errors)\n",
    "        except:\n",
    "            mae = 0\n",
    "        print(f\"{mae:.2f}\", end=',')\n",
    "    print()\n",
    "        \n",
    "\n",
    "print(f\"MAEs for various regions of TorRNA\")\n",
    "for each_angle_name in all_dihedral_angle_names:\n",
    "    print(each_angle_name, end=',')\n",
    "print()\n",
    "\n",
    "print_region_maes(unpaired_errors, region_name=\"Unpaired\")\n",
    "print_region_maes(lone_pairs_errors, region_name=\"Lone Pairs\")\n",
    "print_region_maes(pseudoknots_errors, region_name=\"Pseudoknots\")\n",
    "print_region_maes(multiplets_errors, region_name=\"Multiplets\")\n",
    "print_region_maes(noncanonical_pairs_errors, region_name=\"Non-Canonical Pairs\")\n",
    "print_region_maes(all_canonical_pairs_errors, region_name=\"All Canonical Pairs\")\n",
    "print_region_maes(canonical_nested_pairs_errors, region_name=\"Canonical Nested Pairs\")\n",
    "print_region_maes(hairpin_loop_errors, region_name=\"Hairpin Loops\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Found RNA objects precomputed\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  2%|▏         | 4/172 [00:08<04:21,  1.55s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5LYU_B.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  6%|▌         | 10/172 [00:12<02:17,  1.18it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6O7H_G.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  7%|▋         | 12/172 [00:15<03:09,  1.19s/it]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5WTI_B.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 10%|█         | 18/172 [00:21<02:26,  1.05it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7UZX_G.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 12%|█▏        | 20/172 [00:22<01:57,  1.29it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6BK8_E.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 13%|█▎        | 22/172 [00:23<01:35,  1.58it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//2XDB_G.out due to mismatch of number of residues\n",
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5ZWO_G.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 14%|█▍        | 24/172 [00:25<01:40,  1.47it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1U6B_C.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 27/172 [00:27<02:02,  1.18it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//3WQY_C.out '18'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 20%|██        | 35/172 [00:34<01:23,  1.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1C9S_W.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 24%|██▍       | 41/172 [00:37<01:16,  1.72it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8I45_A.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 27%|██▋       | 47/172 [00:40<01:01,  2.03it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5MWI_A.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 29%|██▉       | 50/172 [00:43<01:24,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//2DU3_D.out '18'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 31%|███▏      | 54/172 [00:46<01:35,  1.23it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5AH5_C.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 34%|███▍      | 59/172 [00:49<01:08,  1.64it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7XPL_G.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 42%|████▏     | 72/172 [00:57<00:52,  1.90it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6NUE_H.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 42%|████▏     | 73/172 [00:58<01:10,  1.40it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//2ZUE_B.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 44%|████▍     | 76/172 [01:00<01:00,  1.58it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7WWV_M.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 47%|████▋     | 81/172 [01:03<00:51,  1.77it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1F7U_B.out '18'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 49%|████▉     | 84/172 [01:06<01:13,  1.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//4RDX_C.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 55%|█████▍    | 94/172 [01:17<01:07,  1.15it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8FZT_A.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 58%|█████▊    | 100/172 [01:18<00:23,  3.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5EEU_W.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 61%|██████    | 105/172 [01:21<00:31,  2.12it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8AW3_1.out '55'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 66%|██████▌   | 113/172 [01:27<00:41,  1.44it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//8BVM_U.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 68%|██████▊   | 117/172 [01:30<00:38,  1.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5OIK_P.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 69%|██████▉   | 119/172 [01:31<00:44,  1.19it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7UOB_T.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 70%|███████   | 121/172 [01:34<00:50,  1.01it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5U30_B.out '75'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 72%|███████▏  | 124/172 [01:35<00:24,  1.97it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7SBA_Z.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 73%|███████▎  | 125/172 [01:36<00:30,  1.54it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7ZHG_5.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 75%|███████▌  | 129/172 [01:40<00:39,  1.08it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//7K16_P.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 83%|████████▎ | 143/172 [01:53<00:27,  1.04it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1MNX_A.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      " 84%|████████▎ | 144/172 [01:54<00:25,  1.09it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1GTN_W.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 85%|████████▍ | 146/172 [01:55<00:18,  1.41it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5LWJ_A.out '35'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 90%|████████▉ | 154/172 [02:01<00:13,  1.33it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5Z71_A.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 91%|█████████ | 156/172 [02:03<00:12,  1.32it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//5ED1_B.out not enough values to unpack (expected 3, got 0)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 92%|█████████▏| 159/172 [02:05<00:11,  1.10it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Unable to correct mappings for /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//1E8O_E.out '15'\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 94%|█████████▍| 162/172 [02:08<00:08,  1.14it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "SKIPPING /home2/sriram.devata/rna_project/rna_structure/data/raw_files/all_torrna_pdbs//6BSI_R.out due to mismatch of number of residues\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 172/172 [02:17<00:00,  1.25it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "MAEs for various regions of SPOT-RNA-1D\n",
      "alpha,beta,gamma,delta,chi,epsilon,zeta,eta,theta,\n",
      "Unpaired (27.77\\%),63.52,29.97,48.91,23.69,32.58,26.79,47.87,50.83,50.05,\n",
      "Lone Pairs (5.99\\%),46.32,23.19,37.55,18.67,23.55,22.73,43.84,27.46,36.57,\n",
      "Pseudoknots (2.71\\%),40.46,18.61,31.47,13.76,14.58,18.80,28.78,25.31,28.30,\n",
      "Multiplets (9.24\\%),53.55,25.33,42.99,18.81,25.29,22.78,42.76,28.56,37.53,\n",
      "Non-Canonical Pairs (25.61\\%),50.10,24.64,41.66,17.91,25.48,22.66,38.76,28.14,34.86,\n",
      "All Canonical Pairs (50.51\\%),32.72,15.66,28.26,11.66,12.15,15.44,16.28,15.48,16.05,\n",
      "Canonical Nested Pairs (46.49\\%),36.56,17.33,31.25,12.41,14.54,16.73,19.47,17.94,18.79,\n",
      "Hairpin Loops (11.72\\%),62.66,27.14,42.92,23.57,28.47,26.55,46.69,51.52,46.73,\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    }
   ],
   "source": [
    "# ------------------- Check SPOT-RNA-1D predictions in various regions -------------------\n",
    "\n",
    "import os\n",
    "if os.path.exists(f\"/ssd_scratch/users/sriram.devata/rna_objects_temp.pkl\"):\n",
    "\twith open(f\"/ssd_scratch/users/sriram.devata/rna_objects_temp.pkl\", \"rb\") as input_file:\n",
    "\t\tprint(f\"Found RNA objects precomputed\")\n",
    "\t\trna_objects = pickle.load(input_file)\n",
    "else:\n",
    "\trna_objects = list()\n",
    "\tprint(f\"Making RNA objects\")\n",
    "\tfor each_pdb_file in tqdm(testing_pdbs):\n",
    "\t\ttry:\n",
    "\t\t\trna_objects.append(RNA(each_pdb_file, calc_rna_fm_embeddings=True, load_dssr_dihedrals=True, load_coords=False))\n",
    "\t\texcept:\n",
    "\t\t\tpass\n",
    "\twith open(f\"/ssd_scratch/users/sriram.devata/rna_objects_temp.pkl\", \"wb\") as output_file:\n",
    "\t\tpickle.dump(rna_objects, output_file)\n",
    "    \n",
    "total_num_nucleotides = 0\n",
    "\n",
    "unpaired_errors = [list() for _ in range(9)]\n",
    "lone_pairs_errors = [list() for _ in range(9)]\n",
    "pseudoknots_errors = [list() for _ in range(9)]\n",
    "multiplets_errors = [list() for _ in range(9)]\n",
    "noncanonical_pairs_errors = [list() for _ in range(9)]\n",
    "all_canonical_pairs_errors = [list() for _ in range(9)]\n",
    "canonical_nested_pairs_errors = [list() for _ in range(9)]\n",
    "hairpin_loop_errors = [list() for _ in range(9)]\n",
    "\n",
    "spot_rna_pred_location = \"/home2/sriram.devata/rna_project/SPOT-RNA-1D/outputs/\"\n",
    "for each_pdb_idx in tqdm(range(len(testing_pdbs))):\n",
    "    \n",
    "    try:\n",
    "        pdb_file_name = testing_pdbs[each_pdb_idx]\n",
    "        rna_object = rna_objects[each_pdb_idx]\n",
    "        gt_angles = rna_object.dssr_torsion_angles\n",
    "        # print(pdb_file_name, len(rna_object.dssr_full_seq))\n",
    "\n",
    "        pdb_code = pdb_file_name.split('/')[-1].replace(\".pdb\", \"\")\n",
    "        spot_rna_pred_file = f\"{spot_rna_pred_location}/{pdb_code}.txt\"\n",
    "\n",
    "        with open(spot_rna_pred_file, \"r\") as f:\n",
    "            spot_rna_pred_lines = f.readlines()\n",
    "\n",
    "        each_rna_predicted_angles = list()\n",
    "        each_rna_gt_angles = list()\n",
    "        # print(pdb_code, gt_angles.shape)\n",
    "        for each_residue_idx in range(len(gt_angles)):\n",
    "            spot_rna_pred_line = spot_rna_pred_lines[each_residue_idx+2]\t# first two lines are column headers and an empty line\n",
    "            cur_res_spot_rna_pred_angles = [float(x) for x in spot_rna_pred_line.strip().split()[2:]]\t# first two columns are number and basename\n",
    "    \n",
    "                        # correcting the order of the angles, SPOT-RNA-1D has a different ordering of angles\n",
    "            index_correction = [0, 1, 2, 3, 6, 4, 5, 7, 8]\n",
    "            cur_res_spot_rna_pred_angles = [cur_res_spot_rna_pred_angles[x] for x in index_correction]\n",
    "\n",
    "            cur_res_gt_angles = make_rad_angle_tensor_to_deg_angle_list(gt_angles[each_residue_idx])\n",
    "\n",
    "            each_rna_predicted_angles.append(cur_res_spot_rna_pred_angles)\n",
    "            each_rna_gt_angles.append(cur_res_gt_angles)\n",
    "\n",
    "\n",
    "        regions_output = get_all_regions(pdb_code, pdbs_path)\n",
    "        if regions_output is not None:\n",
    "            len_seq, all_gt_angles, all_predicted_angles, nuc_num_in_chain_to_zero_index_mapping, unpaired_nums, lone_base_pairs_nums, pseudoknot_nums, multiplet_nums, non_canonical_nums, all_canonical_nums, hairpin_loop_nums, canonical_nested_nums, reshaped_all_base_pairs = regions_output\n",
    "        else:\n",
    "            continue\n",
    "        \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, unpaired_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            unpaired_errors[angle_idx] += all_errors[angle_idx]\n",
    "        \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, lone_base_pairs_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            lone_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "            \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, pseudoknot_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            pseudoknots_errors[angle_idx] += all_errors[angle_idx]\n",
    "            \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, multiplet_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            multiplets_errors[angle_idx] += all_errors[angle_idx]\n",
    "            \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, non_canonical_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            noncanonical_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "        \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, all_canonical_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            all_canonical_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "        \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, canonical_nested_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            canonical_nested_pairs_errors[angle_idx] += all_errors[angle_idx]\n",
    "            \n",
    "        return_maes, all_errors = calculate_mae_for_positions(each_rna_predicted_angles, all_gt_angles, hairpin_loop_nums, prediction_method=\"RNA Transformer Decoder\")\n",
    "        for angle_idx,each_angle_error in enumerate(all_errors):\n",
    "            hairpin_loop_errors[angle_idx] += all_errors[angle_idx]\n",
    "        \n",
    "        total_num_nucleotides += len_seq\n",
    "\n",
    "        if each_pdb_idx > 1000:\n",
    "            break\n",
    "            \n",
    "    except Exception as e:\n",
    "        print(\"Skipping because\", e)\n",
    "        break\n",
    "        pass\n",
    "\n",
    "def print_region_maes(region_errors, region_name=\"Unknown\"):\n",
    "    print(f\"{region_name} ({len(region_errors[0])/total_num_nucleotides*100:.2f}\\%)\", end=',')\n",
    "    for each_angle_idx, each_angle_errors in enumerate(region_errors):\n",
    "        try:\n",
    "            mae = sum(each_angle_errors)/len(each_angle_errors)\n",
    "        except:\n",
    "            mae = 0\n",
    "        print(f\"{mae:.2f}\", end=',')\n",
    "    print()\n",
    "        \n",
    "\n",
    "print(f\"MAEs for various regions of SPOT-RNA-1D\")\n",
    "for each_angle_name in all_dihedral_angle_names:\n",
    "    print(each_angle_name, end=',')\n",
    "print()\n",
    "\n",
    "print_region_maes(unpaired_errors, region_name=\"Unpaired\")\n",
    "print_region_maes(lone_pairs_errors, region_name=\"Lone Pairs\")\n",
    "print_region_maes(pseudoknots_errors, region_name=\"Pseudoknots\")\n",
    "print_region_maes(multiplets_errors, region_name=\"Multiplets\")\n",
    "print_region_maes(noncanonical_pairs_errors, region_name=\"Non-Canonical Pairs\")\n",
    "print_region_maes(all_canonical_pairs_errors, region_name=\"All Canonical Pairs\")\n",
    "print_region_maes(canonical_nested_pairs_errors, region_name=\"Canonical Nested Pairs\")\n",
    "print_region_maes(hairpin_loop_errors, region_name=\"Hairpin Loops\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Region & alpha & beta & gamma & delta & chi & epsilon & zeta & eta & theta &  \\\\\n",
      "\\midrule\n",
      "Unpaired (27.77\\%) & \\shortstack{ 61.56\\\\(63.52) } & \\shortstack{ 29.32\\\\(29.97) } & \\shortstack{ 48.21\\\\(48.91) } & \\shortstack{ 22.45\\\\(23.69) } & \\shortstack{ 32.14\\\\(32.58) } & \\shortstack{ 26.75\\\\(26.79) } & \\shortstack{ 46.76\\\\(47.87) } & \\shortstack{ 48.18\\\\(50.83) } & \\shortstack{ 48.83\\\\(50.05) } & \\\\\n",
      "Lone Pairs (5.99\\%) & \\shortstack{ 44.43\\\\(46.32) } & \\shortstack{ 22.18\\\\(23.19) } & \\shortstack{ 37.18\\\\(37.55) } & \\shortstack{ 16.28\\\\(18.67) } & \\shortstack{ 22.48\\\\(23.55) } & \\shortstack{ 22.96\\\\(\\textbf { 22.73 } ) } & \\shortstack{ 42.93\\\\(43.84) } & \\shortstack{ 25.10\\\\(27.46) } & \\shortstack{ 36.16\\\\(36.57) } & \\\\\n",
      "Pseudoknots (2.71\\%) & \\shortstack{ 38.30\\\\(40.46) } & \\shortstack{ 18.51\\\\(18.61) } & \\shortstack{ 30.23\\\\(31.47) } & \\shortstack{ 11.88\\\\(13.76) } & \\shortstack{ 13.38\\\\(14.58) } & \\shortstack{ 18.73\\\\(18.80) } & \\shortstack{ 27.96\\\\(28.78) } & \\shortstack{ 23.60\\\\(25.31) } & \\shortstack{ 27.61\\\\(28.30) } & \\\\\n",
      "Multiplets (9.24\\%) & \\shortstack{ 50.89\\\\(53.55) } & \\shortstack{ 24.89\\\\(25.33) } & \\shortstack{ 42.53\\\\(42.99) } & \\shortstack{ 17.36\\\\(18.81) } & \\shortstack{ 25.07\\\\(25.29) } & \\shortstack{ 22.88\\\\(\\textbf { 22.78 } ) } & \\shortstack{ 42.60\\\\(42.76) } & \\shortstack{ 27.60\\\\(28.56) } & \\shortstack{ 37.04\\\\(37.53) } & \\\\\n",
      "Non-Canonical Pairs (25.61\\%) & \\shortstack{ 47.24\\\\(50.10) } & \\shortstack{ 24.14\\\\(24.64) } & \\shortstack{ 41.32\\\\(41.66) } & \\shortstack{ 15.58\\\\(17.91) } & \\shortstack{ 24.48\\\\(25.48) } & \\shortstack{ 22.42\\\\(22.66) } & \\shortstack{ 37.94\\\\(38.76) } & \\shortstack{ 25.94\\\\(28.14) } & \\shortstack{ 33.89\\\\(34.86) } & \\\\\n",
      "All Canonical Pairs (50.51\\%) & \\shortstack{ 29.51\\\\(32.72) } & \\shortstack{ 15.06\\\\(15.66) } & \\shortstack{ 27.54\\\\(28.26) } & \\shortstack{ 7.82\\\\(11.66) } & \\shortstack{ 9.37\\\\(12.15) } & \\shortstack{ 13.80\\\\(15.44) } & \\shortstack{ 15.18\\\\(16.28) } & \\shortstack{ 11.10\\\\(15.48) } & \\shortstack{ 14.43\\\\(16.05) } & \\\\\n",
      "Canonical Nested Pairs (46.49\\%) & \\shortstack{ 33.32\\\\(36.56) } & \\shortstack{ 16.89\\\\(17.33) } & \\shortstack{ 30.45\\\\(31.25) } & \\shortstack{ 9.01\\\\(12.41) } & \\shortstack{ 12.43\\\\(14.54) } & \\shortstack{ 15.13\\\\(16.73) } & \\shortstack{ 18.51\\\\(19.47) } & \\shortstack{ 13.95\\\\(17.94) } & \\shortstack{ 17.30\\\\(18.79) } & \\\\\n",
      "Hairpin Loops (11.72\\%) & \\shortstack{ 62.57\\\\(62.66) } & \\shortstack{ 26.55\\\\(27.14) } & \\shortstack{ 42.77\\\\(42.92) } & \\shortstack{ 22.75\\\\(23.57) } & \\shortstack{ 28.52\\\\(\\textbf { 28.47 } ) } & \\shortstack{ 26.57\\\\(\\textbf { 26.55 } ) } & \\shortstack{ 46.18\\\\(46.69) } & \\shortstack{ 48.98\\\\(51.52) } & \\shortstack{ 45.43\\\\(46.73) } & \\\\\n"
     ]
    }
   ],
   "source": [
    "torrna_maes_string=\"\"\"\n",
    "MAEs for various regions of TorRNA\n",
    "alpha,beta,gamma,delta,chi,epsilon,zeta,eta,theta,\n",
    "Unpaired (27.77\\%),61.56,29.32,48.21,22.45,32.14,26.75,46.76,48.18,48.83,\n",
    "Lone Pairs (5.99\\%),44.43,22.18,37.18,16.28,22.48,22.96,42.93,25.10,36.16,\n",
    "Pseudoknots (2.71\\%),38.30,18.51,30.23,11.88,13.38,18.73,27.96,23.60,27.61,\n",
    "Multiplets (9.24\\%),50.89,24.89,42.53,17.36,25.07,22.88,42.60,27.60,37.04,\n",
    "Non-Canonical Pairs (25.61\\%),47.24,24.14,41.32,15.58,24.48,22.42,37.94,25.94,33.89,\n",
    "All Canonical Pairs (50.51\\%),29.51,15.06,27.54,7.82,9.37,13.80,15.18,11.10,14.43,\n",
    "Canonical Nested Pairs (46.49\\%),33.32,16.89,30.45,9.01,12.43,15.13,18.51,13.95,17.30,\n",
    "Hairpin Loops (11.72\\%),62.57,26.55,42.77,22.75,28.52,26.57,46.18,48.98,45.43,\"\"\"\n",
    "\n",
    "spotrna1d_maes_string=\"\"\"\n",
    "MAEs for various regions of SPOT-RNA-1D\n",
    "alpha,beta,gamma,delta,chi,epsilon,zeta,eta,theta,\n",
    "Unpaired (27.77\\%),63.52,29.97,48.91,23.69,32.58,26.79,47.87,50.83,50.05,\n",
    "Lone Pairs (5.99\\%),46.32,23.19,37.55,18.67,23.55,22.73,43.84,27.46,36.57,\n",
    "Pseudoknots (2.71\\%),40.46,18.61,31.47,13.76,14.58,18.80,28.78,25.31,28.30,\n",
    "Multiplets (9.24\\%),53.55,25.33,42.99,18.81,25.29,22.78,42.76,28.56,37.53,\n",
    "Non-Canonical Pairs (25.61\\%),50.10,24.64,41.66,17.91,25.48,22.66,38.76,28.14,34.86,\n",
    "All Canonical Pairs (50.51\\%),32.72,15.66,28.26,11.66,12.15,15.44,16.28,15.48,16.05,\n",
    "Canonical Nested Pairs (46.49\\%),36.56,17.33,31.25,12.41,14.54,16.73,19.47,17.94,18.79,\n",
    "Hairpin Loops (11.72\\%),62.66,27.14,42.92,23.57,28.47,26.55,46.69,51.52,46.73,\"\"\"\n",
    "\n",
    "\n",
    "print(\"Region,alpha,beta,gamma,delta,chi,epsilon,zeta,eta,theta,\".replace(',', ' & '), \"\\\\\\\\\")\n",
    "print(\"\\midrule\")\n",
    "for torrna_line,spotrna1d_line in zip(torrna_maes_string.split('\\n'),spotrna1d_maes_string.split('\\n')):\n",
    "    if '%' not in torrna_line:\n",
    "        continue\n",
    "\n",
    "    torrna_line = torrna_line.split(',')\n",
    "    spotrna1d_line = spotrna1d_line.split(',')\n",
    "    assert torrna_line[0] == spotrna1d_line[0]\n",
    "    \n",
    "    print(f\"{torrna_line[0]}\", end=' & ')\n",
    "    for each_torrna_mae,each_spotrna_mae in zip(torrna_line[1:-1],spotrna1d_line[1:-1]):\n",
    "        # print(f\"\\\\shortstack{{ {each_torrna_mae}\\\\\\\\({each_spotrna_mae}) }}\", end=' & ')\n",
    "        if float(each_torrna_mae) < float(each_spotrna_mae):\n",
    "            print(f\"\\\\shortstack{{ {each_torrna_mae}\\\\\\\\({each_spotrna_mae}) }}\", end=' & ')\n",
    "        else:\n",
    "            print(f\"\\\\shortstack{{ {each_torrna_mae}\\\\\\\\(\\\\textbf {{ {each_spotrna_mae} }} ) }}\", end=' & ')\n",
    "            \n",
    "#         print(f\"{each_torrna_mae}({each_spotrna_mae})\", end=' & ')\n",
    "    print(\"\\\\\\\\\")"
   ]
  },
  {
   "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
}