{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b9515871-d8db-4ccf-8968-88c137134830",
   "metadata": {},
   "source": [
    "Example 2: Detecting topologically non-trivial sequence rearrangements\n",
    "=======================================================================\n",
    "\n",
    "This example demonstrates how GWProt can identify topological permutations in protein sequences, specifically in viral RNA-dependent RNA polymerases (RdRps).\n",
    "\n",
    "Riboviruses are a realm of viruses that include many human pathogens. All riboviruses encode an essential RNA-dependent RNA polymerase (RdRp), which is crucial for viral replication, exhibits a high mutation rate, and often shows very low sequence homology. As a result, structural comparison is key for identification and classification.\n",
    "\n",
    "Viral RdRps contain topologically non-trivial permutations across multiple phyla—proteins where the same motifs or domains appear in different orders along the backbone. RdRps have three active site motifs, labeled A, B, and C, which typically appear in the order ABC but can also appear as CAB. We use 'ABC' and 'CAB' to denote these two types of RdRps."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88614e5c-49fb-4807-bebe-de7cebfd9aef",
   "metadata": {},
   "source": [
    "![Example_Data/Images/ABC_Regions.PNG](Example_Data/Images/ABC_Regions.PNG)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66562dd9-853f-4527-8771-b5a2aeef115d",
   "metadata": {},
   "source": [
    "We first load 10 proteins of each type, with PDB files predicted using AlphaFold."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef0caccd-9f73-4ad6-90ee-585d67e4ec28",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import GWProt.GW_protein\n",
    "\n",
    "dir = 'Example_Data/Permuted Proteins/'\n",
    "file_list = os.listdir(dir)\n",
    "ABC_prots = [GWProt.GW_protein.GW_protein.make_protein_from_pdb(os.path.join(dir , file)) for file in file_list if 'abc' in file]\n",
    "CAB_prots = [GWProt.GW_protein.GW_protein.make_protein_from_pdb(os.path.join(dir , file)) for file in file_list if 'cab' in file]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d3a0d4f4-e2f1-4175-8de1-82141a850146",
   "metadata": {},
   "source": [
    "We then create an `LGD_Comparison` object and compute all pairwise GW distances using local geometric distortion (LGD)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19f2058b-8fe5-49df-93d8-4e0e588d3933",
   "metadata": {},
   "outputs": [],
   "source": [
    "import GWProt.lgd_comparison\n",
    "\n",
    "A = GWProt.lgd_comparison.LGD_Comparison(prot_list=ABC_prots + CAB_prots)\n",
    "A.GW_compute_lgd()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3573dc3a-25c4-4027-a29a-ec4e152edd61",
   "metadata": {},
   "source": [
    "Using UMAP visualization, we observe that proteins can be clustered by permutation type using GW distance, although the separation is weak."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb3a64b0-edda-4f71-8282-5de8c6ef714b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "from umap.umap_ import UMAP\n",
    "\n",
    "reducer = UMAP(metric = 'precomputed')\n",
    "embedding = reducer.fit_transform(A.get_GW_dmat())\n",
    "plt.figure(figsize=(5,5)) \n",
    "plt.scatter(embedding[:10,0] ,embedding[:10,1], color = ['red']*10 , s = 8 )\n",
    "plt.scatter(embedding[10:,0] ,embedding[10:,1], color = ['blue']*10 , s = 8 )\n",
    "plt.legend([\"ABC\",\"CAB\" ])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a7c5191",
   "metadata": {},
   "source": [
    "![Example_Data/Images/Umap3.png](Example_Data/Images/Umap3.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dbf7ea7-9305-4c1f-91f3-5fac8ed40e8a",
   "metadata": {},
   "source": [
    "If we examine the correspondence between two ABC proteins or two CAB proteins, it closely resembles the identity matrix, indicating that the GW alignment is nearly sequential."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a7b43a7-980d-4ea7-9d28-d6a13d051e4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig , axes = plt.subplots(nrows = 1, ncols = 2, figsize=(10,10))\n",
    "fig.tight_layout()\n",
    "\n",
    "# two ABC proteins\n",
    "axes[0].imshow(A.transport_dict[A.name_list[0]][A.name_list[1]], cmap = 'hot_r')\n",
    "axes[0].set_title('ABC to ABC')\n",
    "axes[0].set_xlabel('residue indices')\n",
    "axes[0].set_ylabel('residue indices')\n",
    "\n",
    "\n",
    "#two CAB proteins\n",
    "axes[1].imshow(A.transport_dict[A.name_list[10]][A.name_list[11]], cmap = 'hot_r')\n",
    "axes[1].set_title('CAB to CAB')\n",
    "axes[1].set_xlabel('residue indices')\n",
    "axes[1].set_ylabel('residue indices')\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fef73ae9",
   "metadata": {},
   "source": [
    "![Example_Data/Images/Transport1.png](Example_Data/Images/Transport1.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d9e376c1-0ec8-4178-b1bb-930cf9f39d96",
   "metadata": {},
   "source": [
    "In contrast, the correspondence between an ABC protein and a CAB protein is close to the identity matrix except for a region where the C motifs are aligned."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42af4586-31cc-4509-81b1-384066d67918",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.imshow(A.transport_dict[A.name_list[0]][A.name_list[10]], cmap = 'hot_r', label = 'ABC to CAB')\n",
    "plt.title('ABC to CAB')\n",
    "plt.xlabel('CAB residue indices')\n",
    "plt.ylabel('ABC residue indices')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fefad30e",
   "metadata": {},
   "source": [
    "![Example_Data/Images/Transport2.png](Example_Data/Images/Transport2.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e166a043-6b51-4a44-be36-1e801a455f87",
   "metadata": {},
   "source": [
    "We can visualize the correspondence in PyMOL. The function `pymol_protein_viewer.compare_proteins_in_pymol` uses `weighted_alignment.weighted_RMSD` to superimpose the two proteins by minimizing weighted RMSD, and displays dashes between pairs of aligned residues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c6d5678-7b84-4db5-9a1e-e86090579fe3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import GWProt.pymol_protein_viewer\n",
    "\n",
    "GWProt.pymol_protein_viewer.compare_proteins_in_pymol(\n",
    "    file1 = dir + A.name_list[0] + '.pdb', \n",
    "    file2 = dir + A.name_list[10] + '.pdb',\n",
    "    chain1 ='A', \n",
    "    chain2 ='A', \n",
    "    correspondence=A.transport_dict[A.name_list[0]][A.name_list[10]], \n",
    "    output_file= 'Example_Data/Pymol Outputs/Permutations.pse') "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afc5a2c9-f0ad-461d-b6ce-cd0a8b8125ce",
   "metadata": {},
   "source": [
    "In this screenshot, the two proteins are overlaid: the ABC protein is green, the CAB protein is blue, and redder sections indicate higher local geometric distortion (LGD). The correspondence is shown by yellow lines connecting paired residues."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8ac50532-3a5e-4b27-86ae-e214ce3d980a",
   "metadata": {},
   "source": [
    "![Example_Data/Images/RdRp_ABC_CAB.PNG](Example_Data/Images/RdRp_ABC_CAB.PNG)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3e235a9-2703-4ec3-969c-37066b75cccc",
   "metadata": {},
   "source": [
    "We can use `switch_probabilities.get_switch_probabilities` to identify switched residues, and `switch_probabilities.visualize_switch_probabibilities` for visualization."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73190344-c442-4236-bd75-1a91b6c5c7a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import GWProt.switch_probabilities\n",
    "\n",
    "GWProt.switch_probabilities.visualize_switch_probabibilities(GWProt.switch_probabilities.get_switch_probabilities(A.transport_dict[A.name_list[0]][A.name_list[10]], prot_num=0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d73bf78e",
   "metadata": {},
   "source": [
    "![Example_Data/Images/Triangle.png](Example_Data/Images/Triangle.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e6f5f0c-4a78-47a9-b5a7-50e2ea44439f",
   "metadata": {},
   "source": [
    "The white and light-colored areas represent residue pairs whose order is preserved, while dark rectangles indicate pairs whose order is switched. `prot_num = 0` refers to the residues of the first protein (an ABC protein). The largest dark rectangle shows that residues 211–303 switch with residues 305–318, corresponding to the A–B and C regions, respectively."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cf18fa2-1c6c-4f3a-8801-09477435e369",
   "metadata": {},
   "source": [
    "Using `switch_probabilities.preprocess` and `switch_probabilities.max_rectangle_diagonal`, we can find the size of the largest such rectangle for each protein pair. This can help distinguish permutation types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4be6f685-f369-4cd9-a92d-024108a4b57f",
   "metadata": {},
   "outputs": [],
   "source": [
    "switch_score_mat = np.zeros((20,20))\n",
    "\n",
    "for i in range(20):\n",
    "    for j in range(20):\n",
    "        if i ==j:\n",
    "            continue\n",
    "\n",
    "        T = A.transport_dict[A.name_list[i]][A.name_list[j]]\n",
    "\n",
    "        SP0 = GWProt.switch_probabilities.get_switch_probabilities(T)\n",
    "        SP1 = GWProt.switch_probabilities.preprocess(SP0)\n",
    "        s = GWProt.switch_probabilities.max_rectangle_diagonal(SP1)[0]\n",
    "        switch_score_mat[i,j] = s\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb68e7d9-93dd-492a-bc20-b5e24143a5d8",
   "metadata": {},
   "source": [
    "We can now cluster proteins based on the size of these rectangles, which provides better separation of permutation types."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08bb14ae-d262-4bb5-900e-38ea3f956977",
   "metadata": {},
   "outputs": [],
   "source": [
    "reducer = UMAP(metric = 'precomputed')\n",
    "embedding = reducer.fit_transform(switch_score_mat + switch_score_mat.T)\n",
    "plt.figure(figsize=(5,5)) \n",
    "plt.scatter(embedding[:10,0] ,embedding[:10,1], color = ['red']*10 , s = 8 )\n",
    "plt.scatter(embedding[10:,0] ,embedding[10:,1], color = ['blue']*10 , s = 8 )\n",
    "plt.legend([\"ABC\",\"CAB\" ])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99e03e1b",
   "metadata": {},
   "source": [
    "![Example_Data/Images/Umap4.png](Example_Data/Images/Umap4.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b286952f-06f4-4100-a516-7a0b911e922f",
   "metadata": {},
   "source": [
    "The methods in `switch_probabilities` work well on this dataset because the proteins are morphologically similar in the GW metric. For proteins with greater variation, these techniques may be less effective."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}