project.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Author: Rory Portman <br>\n",
    "Email:&nbsp;&nbsp; portmanrk@gmail.com\n",
    "\n",
    "This notebook forms part of my BSc research project on Binding Free Energy Calculation protocols for Small molecule:RNA Complexes.\n",
    "\n",
    "# BioSimSpace\n",
    "\n",
    "## Molecular Dynamics Simulations and Relative Binding Free Energy Calculations Setup and Analysis for Small Molecule:RNA Complexes\n",
    "\n",
    "This tutorial covers the necessary steps to setup, run and analyse RBFE calculations for Small Molecule:RNA complexes. Included in this tutorial:\n",
    "\n",
    "* Loading, parameterising, solvating and merging the Ligands, the RNA and the Waters.\n",
    "* Running Equilibration and Minimisation for the Small Molecule:RNA system.\n",
    "* Writing scripts to initiate 100 ns molecular dynamics simulations using various simulation engines.\n",
    "* Analysing MD simulations and producing plots for reports.\n",
    "\n",
    "* Loading, parameterising, solvating and merging ligands, the RNA and the waters for RBFE calculations. Includes making merged molecules.\n",
    "* Generating the RBFE input (.pert) files\n",
    "* Analysing the RBFE\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For this tutorial we will be using (ligands and rna from Connely 2019 etc.) \"include introduction for ligands and rna here. reference and provide links\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Firstly let's import all necessary directories below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "062ba90e09714547b576c4725a66c121",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:rdkit:Enabling RDKit 2023.03.3 jupyter extensions\n",
      "INFO:numexpr.utils:Note: NumExpr detected 20 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n"
     ]
    }
   ],
   "source": [
    "#IMPORT DIRECTORIES\n",
    "\n",
    "from get_tutorial import download\n",
    "download()\n",
    "\n",
    "import BioSimSpace as BSS\n",
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from BioSimSpace import _Exceptions\n",
    "import sys\n",
    "import csv"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective 2: Small Molecule:RNA system setup."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before beginning RBFE calculations it is important to determine whether the solvated Small Moleucle:Rna system is stable and how the RNA flexibility affects the ligand binding site.\n",
    "\n",
    "First we must load in the ligands (include name and shorthand name) and the RNA, and parameterise them using existing BioSimSpace Forcefields (include a link?).\n",
    "\n",
    "- Ligand_1: 2-[(dibenzo[b,d]furan-2-yl)oxy]ethan-1-amine\n",
    "- Ligand_2: 2-[(dibenzo[b,d]furan-2-yl)oxy]-N,N-dimethylethan-1-amine\n",
    "- Ligand_3: 2-[(9H-carbazol-3-yl)oxy]-N,N-dimethylethan-1-amine\n",
    "\n",
    "The RNA required the following approach to manually parameterise: \n",
    "\n",
    "\n",
    "The waters were parameterised using the following approach:\n",
    "\n",
    "\n",
    "Prior to loading, ligand 1 is protonated using obabel taking the charge of the ligand into account:\n",
    "\n",
    "* obabel -i pdb lig_1.pdb -o pdb -O lig_1_protonated.pdb -p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ligand = <BioSimSpace.Molecule: number=2, nAtoms=31, nResidues=1>\n",
      "free_ligand = <BioSimSpace.Molecule: number=2, nAtoms=31, nResidues=1>\n",
      "rna =  <BioSimSpace.System: nMolecules=1>\n",
      "water = <BioSimSpace.System: nMolecules=83>\n",
      "rna_bases = <BioSimSpace.System: nMolecules=2>\n"
     ]
    }
   ],
   "source": [
    "#LOAD LIGAND, FREE_LIGAND, RNA AND WATER\n",
    "\n",
    "ligand = BSS.IO.readMolecules(\"input_rory/lig_1_protonated.pdb\")[0]\n",
    "free_ligand = ligand\n",
    "water=BSS.IO.readMolecules(\"input_rory/waters.*\")\n",
    "rna_bases = BSS.IO.readMolecules(\"input_rory/input_rory_3q50/rna.*\")\n",
    "\n",
    "print (\"ligand =\", ligand)\n",
    "print(\"free_ligand =\",free_ligand)\n",
    "print(\"rna = \",rna)\n",
    "print(\"water =\",water)\n",
    "print (\"rna_bases =\", rna_bases)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "To parameterise the ligand we will use the GAFF2 amber forcefield:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "#PARAMATERISE LIGAND, FREE_LIGAND AND WATER\n",
    "\n",
    "gaff_free_ligand = BSS.Parameters.gaff2(free_ligand).getMolecule()\n",
    "gaff_ligand = BSS.Parameters.gaff2(ligand).getMolecule()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we combine the parameterised RNA, ligand and waters into one system called `molecules`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "74497cb2fe144d798810615b8ad6308e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "ThemeManager()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "caf89d3d070e4ce696624dd89883a179",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget(gui_style='ngl')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#COMBINE RNA, WATER AND LIGAND\n",
    "\n",
    "molecules= rna_bases + gaff_ligand + water\n",
    "view=BSS.Notebook.View(rna_bases)\n",
    "view.system()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is useful to check our system using `BSS.Notebook.View`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "5d7d3a8095524bd4b09420e7392bd2b3",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget(gui_style='ngl')"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "view_2= BSS.Notebook.View(\"input_rory/lig_protonated.pdb\")\n",
    "view_2.system()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After creating a Small Molecule:RNA system we must solvate using `BSS.Solvent`. In this case we will use the `tip3p` solvating model and we will automatically calculate the necessary solvent box size using `.getAxisAlignedBoundingBox()`. It is important to specify the counter ion concentration."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#DETERMINE BOX SIZE AND SOLVATE COMBINED SYSTEM (CHECK COUNTER IONS)\n",
    "\n",
    "box_min, box_max = rna.getAxisAlignedBoundingBox()\n",
    "box_size = [y - x for x, y in zip(box_min, box_max)]\n",
    "padding = 15 * BSS.Units.Length.angstrom\n",
    "box_length = max(box_size) + 2 * padding\n",
    "\n",
    "molecules_sol = BSS.Solvent.tip3p(molecule= molecules, box=3 * [box_length], ion_conc = 0.15)\n",
    "search = molecules_sol.search(\"not mols with atomidx 2\")\n",
    "for ion in search.atoms():\n",
    "    print(f\"element = {ion.element()}, charge = {ion.charge()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And the same for the free ligand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#DETERMINE BOX SIZE AND SOLVATE FREE LIGAND\n",
    "\n",
    "free_box_min, free_box_max = free_ligand.getAxisAlignedBoundingBox()\n",
    "free_box_size = [y - x for x, y in zip(free_box_min, free_box_max)]\n",
    "free_padding = 15 * BSS.Units.Length.angstrom\n",
    "free_box_length = max(free_box_size) + 2 * free_padding\n",
    "print(free_box_length)\n",
    "\n",
    "free_molecules_sol = BSS.Solvent.tip3p(molecule= gaff_free_ligand, box=3 * [free_box_length], ion_conc = 0.15)\n",
    "search = free_molecules_sol.search(\"not mols with atomidx 2\")\n",
    "for ion in search.atoms():\n",
    "    print(f\"element = {ion.element()}, charge = {ion.charge()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective 3a: Minimisation and Equilibration of Small Molecule:RNA system."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a solvated Small Molecule:RNA system we need to ensure that the starting structure is at a local energy minimum prior to beginning the 100 ns MD simulations. \n",
    "\n",
    "To do this we will perform equilibration and minimisation of the ligand_1:RNA system.\n",
    "\n",
    "The following approach for minimisation and equilibration is adapted from: https://github.com/michellab/BioSimSpaceTutorials/blob/69e2a9f97b2d2f5ae21c5c5b45a2992b2879621c/03_fep/execution_model/scripts/BSSligprep.py#L225 \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "### Settings.\n",
    "minim_steps = 250\n",
    "runtime_short_nvt = 5 # ps\n",
    "runtime_nvt = 50 # ps RESET to 50 after TESTING ! \n",
    "runtime_npt = 200 # ps RESET to 200 after TESTING !\n",
    "\n",
    "### preamble. tmp_dir should at some point be derived using os.environ.get(\"\")\n",
    "\n",
    "tmp_dir = \"temp\"\n",
    "\n",
    "amber_home = \"/usr/local/amber22\"\n",
    "pmemd_path = amber_home + \"/bin/pmemd.cuda\" "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "lig_p_solvated is = <BioSimSpace.System: nMolecules=2154>\n",
      "original is = <BioSimSpace.System: nMolecules=2154>\n"
     ]
    }
   ],
   "source": [
    "#################\n",
    "lig_p_solvated = free_molecules_sol\n",
    "system_solvated = molecules_sol\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "### Minimise and equilibrate our systems to make the ligand relax inside the pocket.\n",
    "\n",
    "lig_p_solvated = BSS.IO.readMolecules([f\"{tmp_dir}/free_lig_s.prm7\", f\"{tmp_dir}/free_lig_s.rst7\"])\n",
    "system_solvated = BSS.IO.readMolecules([f\"{tmp_dir}/mol_sys_s.prm7\", f\"{tmp_dir}/mol_sys_s.rst7\"])\n",
    "\n",
    "print (\"lig_p_solvated is =\", lig_p_solvated)\n",
    "print (\"original is =\", free_molecules_sol)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "#### Working on solvated ligand.\n",
      "Minimising in 250 steps..\n",
      "PMEMD NVT equilibration for 5 ps while restraining all non-solvent atoms..\n",
      "PMEMD NVT equilibration for 50 ps without restraints..\n",
      "PMEMD NPT equilibration for 200 ps while restraining non-solvent heavy atoms..\n",
      "PMEMD NPT equilibration for 200 ps without restraints..\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "['/home/2020/home/rory/project/ProjectDirectory/temp/eq4_free_lig_s.prm7',\n",
       " '/home/2020/home/rory/project/ProjectDirectory/temp/eq4_free_lig_s.rst7']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "\n",
    "def runProcess(system, protocol, pmemd=False):\n",
    "        \"\"\"\n",
    "        Given a solvated system (BSS object) and BSS protocol, run a process workflow with either \n",
    "        Sander (CPU) or pmemd.cuda (GPU). NPT is typically done with GPU to save computing time.\n",
    "        Returns the processed system.\n",
    "        \"\"\"\n",
    "        # Create the process passing a working directory.\n",
    "        if not pmemd:\n",
    "            process = BSS.Process.Amber(system, protocol)\n",
    "        elif pmemd:\n",
    "            process = BSS.Process.Amber(system, protocol, exe=pmemd_path)\n",
    "\n",
    "        # Start the process.\n",
    "        process.start()\n",
    "\n",
    "        # Wait for the process to exit.\n",
    "        process.wait()\n",
    "\n",
    "        # Check for errors.\n",
    "        if process.isError():\n",
    "            print(process.stdout())\n",
    "            print(process.stderr())\n",
    "            raise _Exceptions.ThirdPartyError(\"The process exited with an error!\")\n",
    "\n",
    "        # If it worked, try to get the system. No need to block, since it's already finished.\n",
    "        system = process.getSystem()\n",
    "\n",
    "        return system\n",
    "\n",
    "\n",
    "############# first minimise/equilibrate the solvated ligand.\n",
    "print(\"\\n#### Working on solvated ligand.\")\n",
    "\n",
    "print(f\"Minimising in {minim_steps} steps..\")\n",
    "protocol = BSS.Protocol.Minimisation(steps=minim_steps)\n",
    "minimised = runProcess(lig_p_solvated, protocol)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/min_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_short_nvt} ps while restraining all non-solvent atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_short_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_start=0*BSS.Units.Temperature.kelvin, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"all\"\n",
    "                                )\n",
    "equil1 = runProcess(minimised, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq1_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "\n",
    "equil2 = runProcess(equil1, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq2_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps while restraining non-solvent heavy atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"heavy\",\n",
    "                                )\n",
    "equil3 = runProcess(equil2, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq3_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "lig_equil_fin = runProcess(equil3, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq4_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "#### Working on solvated ligand+protein.\n",
      "Minimising in 250 steps..\n",
      "PMEMD NVT equilibration for 5 ps while restraining all non-solvent atoms..\n",
      "PMEMD NVT equilibration for 50 ps while restraining all backbone atoms..\n",
      "PMEMD NVT equilibration for 50 ps without restraints..\n",
      "PMEMD NPT equilibration for 200 ps while restraining non-solvent heavy atoms..\n",
      "PMEMD NPT equilibration for 200 ps without restraints..\n",
      "Saving solvated/equilibrated systems.\n",
      "\n",
      " Ligand:\n",
      "<BioSimSpace.System: nMolecules=2154>\n",
      "\n",
      " Ligand + protein:\n",
      "<BioSimSpace.System: nMolecules=10000>\n",
      "First 20 molecules in ligand + protein system:\n",
      "<BioSimSpace.Molecule: number=26899, nAtoms=1062, nResidues=33>\n",
      "<BioSimSpace.Molecule: number=36240, nAtoms=31, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36241, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36242, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36243, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36244, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36245, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36246, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36247, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36248, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36257, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36259, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36277, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36285, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36295, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36306, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36315, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36326, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36338, nAtoms=3, nResidues=1>\n",
      "<BioSimSpace.Molecule: number=36324, nAtoms=3, nResidues=1>\n",
      "Done.\n"
     ]
    }
   ],
   "source": [
    "############# repeat for ligand + protein system. Include extra restrain=\"backbone\" step.\n",
    "print(\"\\n#### Working on solvated ligand+protein.\")\n",
    "print(f\"Minimising in {minim_steps} steps..\")\n",
    "protocol = BSS.Protocol.Minimisation(steps=minim_steps)\n",
    "minimised = runProcess(system_solvated, protocol)\n",
    "\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_short_nvt} ps while restraining all non-solvent atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_short_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_start=0*BSS.Units.Temperature.kelvin, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"all\"\n",
    "                                )\n",
    "equil1 = runProcess(minimised, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq1_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps while restraining all backbone atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature=300*BSS.Units.Temperature.kelvin, \n",
    "                                #restraint=\"backbone\"\n",
    "                                )\n",
    "equil2 = runProcess(equil1, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq2_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "\n",
    "equil3 = runProcess(equil2, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq3_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps while restraining non-solvent heavy atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"heavy\",\n",
    "                                )\n",
    "equil4 = runProcess(equil3, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq4_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "sys_equil_fin = runProcess(equil4, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq5_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "# finally, save last snapshot of both equilibrated objects.\n",
    "os.system(\"mkdir -p prep/ligands\")\n",
    "os.system(\"mkdir -p prep/protein\")\n",
    "\n",
    "print(\"Saving solvated/equilibrated systems.\")\n",
    "print(\"\\n Ligand:\")\n",
    "print(lig_equil_fin)\n",
    "BSS.IO.saveMolecules(f\"prep_bases/ligands/free_lig_equil_solv\", lig_equil_fin, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(\"\\n Ligand + protein:\")\n",
    "print(sys_equil_fin)\n",
    "BSS.IO.saveMolecules(f\"prep_bases/protein/mol_sys_equil_solv\", sys_equil_fin, [\"PRM7\", \"RST7\"])\n",
    "print(\"First 20 molecules in ligand + protein system:\")\n",
    "for mol in sys_equil_fin.getMolecules()[:20]:\n",
    "    print(mol)\n",
    "print(\"Done.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective 3b: Run MD simulation using OpenMM, pmemd.cuda and gromacs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To run a molecular dynamics simulation of the ligand (ligand_1) in the binding site of the RNA (rna_bases) we will make use of the BSS.Protocol and BSS.Process functions. 100 ns MD simulations of the ligand_1:rna_bases system  were run using AMBER, GROMACS and OpenMM. The MD simulations are initiated using the following [Python Script](/home/rory/project/ProjectDirectory/project_scripts_2/amber_run_bases.py) and queued using the following [Slurm Script](/home/rory/project/ProjectDirectory/project_scripts_2/amber_run_bases.sh). An example python script is shopwn below:\n",
    "\n",
    "```python\n",
    "import BioSimSpace as BSS\n",
    "import os\n",
    "\n",
    "print(\"Loading system...\")\n",
    "system= BSS.IO.readMolecules(\"../prep/protein/mol_sys_equil_solv.*\")\n",
    "protocol=BSS.Protocol.Production(runtime= 100 * BSS.Units.Time.nanosecond, report_interval=10_000)\n",
    "process = BSS.Process.Amber(system, protocol,exe = \"/usr/local/amber22/bin/pmemd.cuda\", work_dir=\"output_amber\")\n",
    "process.start()\n",
    "process.wait()\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective 3 Analysis: Plot RMSD and RMSF values for the trajectories."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "RMSD Plots: for nucleic \"nucleic\", for backbone \"nucleic or resid 13 or resid 14\", for ligand \"resname LIG\", for all \"nucleic or resid 13 or resid 14 or resname LIG\"."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rmsd_plot (topology, trajectory, selection, name):\n",
    "    import MDAnalysis as mda\n",
    "    from MDAnalysis.analysis import rms, align\n",
    "    import matplotlib.pyplot as plt\n",
    "    import numpy as np\n",
    "\n",
    "    # Load the trajectory using DCDReader\n",
    "    u = mda.Universe(topology, trajectory)\n",
    "\n",
    "    # Select backbone atoms\n",
    "    backbone_1 = selection\n",
    "\n",
    "    # Calculate RMSD\n",
    "    R = rms.RMSD(u, u, select= backbone_1)\n",
    "    R.run()\n",
    "\n",
    "    # Access RMSD data\n",
    "    results = R.results.rmsd.T\n",
    "    rmsds= results[2]\n",
    "    time= results[1]\n",
    "    fig, ax = plt.subplots()\n",
    "    ax.plot(time, rmsds)\n",
    "    ax.set_xlabel(\"Time / ps\")\n",
    "    ax.set_ylabel(\"RMSD / $\\\\AA$\")\n",
    "    ax.grid(True)\n",
    "    plt.savefig(name)\n",
    "\n",
    "rmsd_plot ('openmm.parm7', 'thinned_openmm.nc', 'nucleic', 'rmsd_openmm_nucleic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The most useful analysis of the RNA system is to plot the RMSF values for the individual residues in the RNA backbone to determine the positions of flexibility in the system and hence determine any potential effects on the binding site. To do this we can use the following rmsf_plot function for all three simulations (AMBER, GROMACS and OpenMM):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def rmsf_plot (topology_1, trajectory_1, selection_1, name_1):\n",
    "     u_1 = mda.Universe(topology_1, trajectory_1)\n",
    "     average = align.AverageStructure(u_1,u_1, select= selection_1, ref_frame=0).run()\n",
    "     ref = average.results.universe\n",
    "\n",
    "     aligner = align.AlignTraj(u_1, ref, select= selection_1, in_memory=True).run()\n",
    "     rna = u_1.select_atoms(selection_1)\n",
    "     R = rms.RMSF(rna).run()\n",
    "     fig, ax = plt.subplots()\n",
    "     ax.plot(R.rmsf)\n",
    "     ax.set_xlabel(\"Relative Atom Index\")\n",
    "     ax.set_ylabel(\"RMSF / $\\\\AA$\")\n",
    "     tick_step = np.arange(0,1001, 100)\n",
    "     ax.set_xticks(tick_step)\n",
    "     plt.savefig(name)\n",
    "\n",
    "rmsf_plot ('openmm.parm7', 'thinned_openmm.nc', 'nucleic', 'rmsf_openmm_nucleic')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Objective 4: RBFE calculations for compounds 1-3 using SOMD and AMBER"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4a: Set up of RBFE calculations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import libraries\n",
    "import BioSimSpace as BSS\n",
    "import os\n",
    "import glob\n",
    "import csv\n",
    "import numpy as np\n",
    "import networkx as nx\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As ligand 1 is already in the binding site following the molecular dynamics simulations from Objective 3 it can be used to create merged molecules with ligands 2 and 3 respectively and complete the alchemical transformations necessary to compute Relative Binding Free Energy. \n",
    "- Ligand 1 in the binding site alchemically transformed into Ligand 2 \n",
    "- Ligand 1 in the binding site alchemically transformed into Ligand 3\n",
    "\n",
    "It is important to calculate the RBFE in both directions e.g.\n",
    "- Ligand 1 in the binding site to Ligand 2 \n",
    "- Ligand 2 in the binding site to Ligand 1 \n",
    "\n",
    "As such we need to create ligand:rna systems for ligand 2 and lind 3 similar to ligand 1, and run minimisation and equilibration. There is no need to run MD simulations for ligand 2 and ligand 3 as the ligands are chemically similar to ligand 1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below we have loaded the RNA system and Ligand 1 from the MD simulations along with Ligand 2, Ligand 3, and the parameterised waters. Ligands 2 and 3 required protonation prior to loading using obabel in the command line. An example of this is:\n",
    "\n",
    "* obabel -i pdb lig_2.pdb -o pdb -O lig_2_protonated.pdb -p"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "system = BSS.IO.readMolecules(\"../prep_bases/protein/mol_sys_equil_solv.*\")\n",
    "ligand_1_aligned = system[1]\n",
    "\n",
    "pre_ligand_1 = BSS.IO.readMolecules(\"protonated_1.pdb\")[0]\n",
    "pre_ligand_2 = BSS.IO.readMolecules(\"protonated_2.pdb\")[0]\n",
    "pre_ligand_3 = BSS.IO.readMolecules(\"protonated_3.pdb\")[0]\n",
    "ligand_1 = BSS.Parameters.gaff2(pre_ligand_1).getMolecule()\n",
    "ligand_2 = BSS.Parameters.gaff2(pre_ligand_2).getMolecule()\n",
    "ligand_3 = BSS.Parameters.gaff2(pre_ligand_3).getMolecule()\n",
    "\n",
    "waters = BSS.IO.readMolecules(\"../input_rory/waters.*\")\n",
    "rna_bases = system[0]\n",
    "rna_bases_test = BSS.IO.readMolecules(\"/home/rory/project/ProjectDirectory/input_rory/input_rory_3q50/rna.*\")\n",
    "\n",
    "# Print the new variable\n",
    "print('system =', system)\n",
    "\n",
    "print ('rna =', rna_bases)\n",
    "print ('ligand -',ligand_1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To make sure Ligand 2 and Ligand 3 are in the binding sites of their respective Ligand:RNA systems we can align them to Ligand 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mapping_3_1 = BSS.Align.matchAtoms(ligand_1_aligned, ligand_3)\n",
    "inv_mapping = {v: k for k, v in mapping_3_1.items()}\n",
    "\n",
    "# Align 3\n",
    "ligand_3_aligned = BSS.Align.rmsdAlign(ligand_3, ligand_1_aligned, inv_mapping)\n",
    "\n",
    "mapping_2_1 = BSS.Align.matchAtoms(ligand_1_aligned, ligand_2)\n",
    "inv_mapping_2 = {v: k for k, v in mapping_2_1.items()}\n",
    "\n",
    "#Align 2\n",
    "ligand_2_aligned = BSS.Align.rmsdAlign(ligand_2, ligand_1_aligned, inv_mapping_2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can equilibrate and minimise Ligand 2 in the binding site."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol_2 = rna_bases + ligand_2_aligned + waters\n",
    "lig_2 = ligand_2_aligned + waters \n",
    "\n",
    "#DETERMINE BOX SIZE AND SOLVATE COMBINED SYSTEM (CHECK COUNTER IONS)\n",
    "\n",
    "box_min, box_max = rna_bases.getAxisAlignedBoundingBox()\n",
    "box_size = [y - x for x, y in zip(box_min, box_max)]\n",
    "padding = 15 * BSS.Units.Length.angstrom\n",
    "box_length = max(box_size) + 2 * padding\n",
    "\n",
    "molecules_sol = BSS.Solvent.tip3p(molecule= mol_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "search = molecules_sol.search(\"not mols with atomidx 2\")\n",
    "for ion in search.atoms():\n",
    "    print(f\"element = {ion.element()}, charge = {ion.charge()}\")\n",
    "    \n",
    "    #DETERMINE BOX SIZE AND SOLVATE COMBINED SYSTEM (CHECK COUNTER IONS)\n",
    "\n",
    "box_min, box_max = rna_bases.getAxisAlignedBoundingBox()\n",
    "box_size = [y - x for x, y in zip(box_min, box_max)]\n",
    "padding = 15 * BSS.Units.Length.angstrom\n",
    "box_length = max(box_size) + 2 * padding\n",
    "\n",
    "free_molecules_sol = BSS.Solvent.tip3p(molecule= lig_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "search = molecules_sol.search(\"not mols with atomidx 2\")\n",
    "for ion in search.atoms():\n",
    "    print(f\"element = {ion.element()}, charge = {ion.charge()}\")\n",
    "    \n",
    "    ### Settings.\n",
    "minim_steps = 250\n",
    "runtime_short_nvt = 5 # ps\n",
    "runtime_nvt = 50 # ps RESET to 50 after TESTING ! \n",
    "runtime_npt = 200 # ps RESET to 200 after TESTING !\n",
    "\n",
    "### preamble. tmp_dir should at some point be derived using os.environ.get(\"\")\n",
    "\n",
    "tmp_dir = \"temp_2\"\n",
    "\n",
    "amber_home = \"/usr/local/amber22\"\n",
    "pmemd_path = amber_home + \"/bin/pmemd.cuda\" \n",
    "\n",
    "#################\n",
    "lig_p_solvated = free_molecules_sol\n",
    "system_solvated = molecules_sol\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "### Minimise and equilibrate our systems to make the ligand relax inside the pocket.\n",
    "\n",
    "lig_p_solvated = BSS.IO.readMolecules([f\"{tmp_dir}/free_lig_s.prm7\", f\"{tmp_dir}/free_lig_s.rst7\"])\n",
    "system_solvated = BSS.IO.readMolecules([f\"{tmp_dir}/mol_sys_s.prm7\", f\"{tmp_dir}/mol_sys_s.rst7\"])\n",
    "\n",
    "print (\"lig_p_solvated is =\", lig_p_solvated)\n",
    "print (\"original is =\", free_molecules_sol)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def runProcess(system, protocol, pmemd=False):\n",
    "        \"\"\"\n",
    "        Given a solvated system (BSS object) and BSS protocol, run a process workflow with either \n",
    "        Sander (CPU) or pmemd.cuda (GPU). NPT is typically done with GPU to save computing time.\n",
    "        Returns the processed system.\n",
    "        \"\"\"\n",
    "        # Create the process passing a working directory.\n",
    "        if not pmemd:\n",
    "            process = BSS.Process.Amber(system, protocol)\n",
    "        elif pmemd:\n",
    "            process = BSS.Process.Amber(system, protocol, exe=pmemd_path)\n",
    "\n",
    "        # Start the process.\n",
    "        process.start()\n",
    "\n",
    "        # Wait for the process to exit.\n",
    "        process.wait()\n",
    "\n",
    "        # Check for errors.\n",
    "        if process.isError():\n",
    "            print(process.stdout())\n",
    "            print(process.stderr())\n",
    "            raise _Exceptions.ThirdPartyError(\"The process exited with an error!\")\n",
    "\n",
    "        # If it worked, try to get the system. No need to block, since it's already finished.\n",
    "        system = process.getSystem()\n",
    "\n",
    "        return system\n",
    "\n",
    "\n",
    "############# first minimise/equilibrate the solvated ligand.\n",
    "print(\"\\n#### Working on solvated ligand.\")\n",
    "\n",
    "print(f\"Minimising in {minim_steps} steps..\")\n",
    "protocol = BSS.Protocol.Minimisation(steps=minim_steps)\n",
    "minimised = runProcess(lig_p_solvated, protocol)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/min_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_short_nvt} ps while restraining all non-solvent atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_short_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_start=0*BSS.Units.Temperature.kelvin, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"all\"\n",
    "                                )\n",
    "equil1 = runProcess(minimised, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq1_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "\n",
    "equil2 = runProcess(equil1, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq2_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps while restraining non-solvent heavy atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"heavy\",\n",
    "                                )\n",
    "equil3 = runProcess(equil2, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq3_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "lig_equil_fin = runProcess(equil3, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq4_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "############# repeat for ligand + protein system. Include extra restrain=\"backbone\" step.\n",
    "print(\"\\n#### Working on solvated ligand+protein.\")\n",
    "print(f\"Minimising in {minim_steps} steps..\")\n",
    "protocol = BSS.Protocol.Minimisation(steps=minim_steps)\n",
    "minimised = runProcess(system_solvated, protocol)\n",
    "\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_short_nvt} ps while restraining all non-solvent atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_short_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_start=0*BSS.Units.Temperature.kelvin, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"all\"\n",
    "                                )\n",
    "equil1 = runProcess(minimised, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq1_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps while restraining all backbone atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature=300*BSS.Units.Temperature.kelvin, \n",
    "                                #restraint=\"backbone\"\n",
    "                                )\n",
    "equil2 = runProcess(equil1, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq2_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "\n",
    "equil3 = runProcess(equil2, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq3_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps while restraining non-solvent heavy atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"heavy\",\n",
    "                                )\n",
    "equil4 = runProcess(equil3, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq4_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "sys_equil_fin = runProcess(equil4, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq5_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "# finally, save last snapshot of both equilibrated objects.\n",
    "os.system(\"mkdir -p prep/ligands\")\n",
    "os.system(\"mkdir -p prep/protein\")\n",
    "\n",
    "print(\"Saving solvated/equilibrated systems.\")\n",
    "print(\"\\n Ligand:\")\n",
    "print(lig_equil_fin)\n",
    "BSS.IO.saveMolecules(f\"prep_bases_2/ligands_2/free_lig_equil_solv\", lig_equil_fin, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(\"\\n Ligand + protein:\")\n",
    "print(sys_equil_fin)\n",
    "BSS.IO.saveMolecules(f\"prep_bases_2/protein_2/mol_sys_equil_solv\", sys_equil_fin, [\"PRM7\", \"RST7\"])\n",
    "print(\"First 20 molecules in ligand + protein system:\")\n",
    "for mol in sys_equil_fin.getMolecules()[:20]:\n",
    "    print(mol)\n",
    "print(\"Done.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can equilibrate and minimise Ligand 3 in the binding site."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mol_2 = rna_bases + ligand_3_aligned + waters\n",
    "lig_2 = ligand_3 + waters \n",
    "\n",
    "#DETERMINE BOX SIZE AND SOLVATE COMBINED SYSTEM (CHECK COUNTER IONS)\n",
    "\n",
    "box_min, box_max = rna_bases.getAxisAlignedBoundingBox()\n",
    "box_size = [y - x for x, y in zip(box_min, box_max)]\n",
    "padding = 15 * BSS.Units.Length.angstrom\n",
    "box_length = max(box_size) + 2 * padding\n",
    "\n",
    "molecules_sol = BSS.Solvent.tip3p(molecule= mol_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "search = molecules_sol.search(\"not mols with atomidx 2\")\n",
    "for ion in search.atoms():\n",
    "    print(f\"element = {ion.element()}, charge = {ion.charge()}\")\n",
    "    \n",
    "    #DETERMINE BOX SIZE AND SOLVATE COMBINED SYSTEM (CHECK COUNTER IONS)\n",
    "\n",
    "box_min, box_max = rna_bases.getAxisAlignedBoundingBox()\n",
    "box_size = [y - x for x, y in zip(box_min, box_max)]\n",
    "padding = 15 * BSS.Units.Length.angstrom\n",
    "box_length = max(box_size) + 2 * padding\n",
    "\n",
    "free_molecules_sol = BSS.Solvent.tip3p(molecule= lig_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "search = molecules_sol.search(\"not mols with atomidx 2\")\n",
    "for ion in search.atoms():\n",
    "    print(f\"element = {ion.element()}, charge = {ion.charge()}\")\n",
    "    \n",
    "    ### Settings.\n",
    "minim_steps = 250\n",
    "runtime_short_nvt = 5 # ps\n",
    "runtime_nvt = 50 # ps RESET to 50 after TESTING ! \n",
    "runtime_npt = 200 # ps RESET to 200 after TESTING !\n",
    "\n",
    "### preamble. tmp_dir should at some point be derived using os.environ.get(\"\")\n",
    "\n",
    "tmp_dir = \"temp_3\"\n",
    "\n",
    "amber_home = \"/usr/local/amber22\"\n",
    "pmemd_path = amber_home + \"/bin/pmemd.cuda\" \n",
    "\n",
    "#################\n",
    "lig_p_solvated = free_molecules_sol\n",
    "system_solvated = molecules_sol\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "### Minimise and equilibrate our systems to make the ligand relax inside the pocket.\n",
    "\n",
    "lig_p_solvated = BSS.IO.readMolecules([f\"{tmp_dir}/free_lig_s.prm7\", f\"{tmp_dir}/free_lig_s.rst7\"])\n",
    "system_solvated = BSS.IO.readMolecules([f\"{tmp_dir}/mol_sys_s.prm7\", f\"{tmp_dir}/mol_sys_s.rst7\"])\n",
    "\n",
    "print (\"lig_p_solvated is =\", lig_p_solvated)\n",
    "print (\"original is =\", free_molecules_sol)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def runProcess(system, protocol, pmemd=False):\n",
    "        \"\"\"\n",
    "        Given a solvated system (BSS object) and BSS protocol, run a process workflow with either \n",
    "        Sander (CPU) or pmemd.cuda (GPU). NPT is typically done with GPU to save computing time.\n",
    "        Returns the processed system.\n",
    "        \"\"\"\n",
    "        # Create the process passing a working directory.\n",
    "        if not pmemd:\n",
    "            process = BSS.Process.Amber(system, protocol)\n",
    "        elif pmemd:\n",
    "            process = BSS.Process.Amber(system, protocol, exe=pmemd_path)\n",
    "\n",
    "        # Start the process.\n",
    "        process.start()\n",
    "\n",
    "        # Wait for the process to exit.\n",
    "        process.wait()\n",
    "\n",
    "        # Check for errors.\n",
    "        if process.isError():\n",
    "            print(process.stdout())\n",
    "            print(process.stderr())\n",
    "            raise _Exceptions.ThirdPartyError(\"The process exited with an error!\")\n",
    "\n",
    "        # If it worked, try to get the system. No need to block, since it's already finished.\n",
    "        system = process.getSystem()\n",
    "\n",
    "        return system\n",
    "\n",
    "\n",
    "############# first minimise/equilibrate the solvated ligand.\n",
    "print(\"\\n#### Working on solvated ligand.\")\n",
    "\n",
    "print(f\"Minimising in {minim_steps} steps..\")\n",
    "protocol = BSS.Protocol.Minimisation(steps=minim_steps)\n",
    "minimised = runProcess(lig_p_solvated, protocol)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/min_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_short_nvt} ps while restraining all non-solvent atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_short_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_start=0*BSS.Units.Temperature.kelvin, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"all\"\n",
    "                                )\n",
    "equil1 = runProcess(minimised, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq1_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "\n",
    "equil2 = runProcess(equil1, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq2_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps while restraining non-solvent heavy atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"heavy\",\n",
    "                                )\n",
    "equil3 = runProcess(equil2, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq3_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "lig_equil_fin = runProcess(equil3, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq4_free_lig_s\", lig_p_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "############# repeat for ligand + protein system. Include extra restrain=\"backbone\" step.\n",
    "print(\"\\n#### Working on solvated ligand+protein.\")\n",
    "print(f\"Minimising in {minim_steps} steps..\")\n",
    "protocol = BSS.Protocol.Minimisation(steps=minim_steps)\n",
    "minimised = runProcess(system_solvated, protocol)\n",
    "\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_short_nvt} ps while restraining all non-solvent atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_short_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_start=0*BSS.Units.Temperature.kelvin, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"all\"\n",
    "                                )\n",
    "equil1 = runProcess(minimised, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq1_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps while restraining all backbone atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature=300*BSS.Units.Temperature.kelvin, \n",
    "                                #restraint=\"backbone\"\n",
    "                                )\n",
    "equil2 = runProcess(equil1, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq2_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NVT equilibration for {runtime_nvt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_nvt*BSS.Units.Time.picosecond, \n",
    "                                temperature_end=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "\n",
    "equil3 = runProcess(equil2, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq3_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps while restraining non-solvent heavy atoms..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                restraint=\"heavy\",\n",
    "                                )\n",
    "equil4 = runProcess(equil3, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq4_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(f\"PMEMD NPT equilibration for {runtime_npt} ps without restraints..\")\n",
    "protocol = BSS.Protocol.Equilibration(\n",
    "                                runtime=runtime_npt*BSS.Units.Time.picosecond, \n",
    "                                pressure=1*BSS.Units.Pressure.atm,\n",
    "                                temperature=300*BSS.Units.Temperature.kelvin,\n",
    "                                )\n",
    "sys_equil_fin = runProcess(equil4, protocol, pmemd=True)\n",
    "\n",
    "BSS.IO.saveMolecules(f\"{tmp_dir}/eq5_mol_sys_s\", system_solvated, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "# finally, save last snapshot of both equilibrated objects.\n",
    "os.system(\"mkdir -p prep/ligands\")\n",
    "os.system(\"mkdir -p prep/protein\")\n",
    "\n",
    "print(\"Saving solvated/equilibrated systems.\")\n",
    "print(\"\\n Ligand:\")\n",
    "print(lig_equil_fin)\n",
    "BSS.IO.saveMolecules(f\"prep_bases_3/ligands_3/free_lig_equil_solv\", lig_equil_fin, [\"PRM7\", \"RST7\"])\n",
    "\n",
    "print(\"\\n Ligand + protein:\")\n",
    "print(sys_equil_fin)\n",
    "BSS.IO.saveMolecules(f\"prep_bases_3/protein_3/mol_sys_equil_solv\", sys_equil_fin, [\"PRM7\", \"RST7\"])\n",
    "print(\"First 20 molecules in ligand + protein system:\")\n",
    "for mol in sys_equil_fin.getMolecules()[:20]:\n",
    "    print(mol)\n",
    "print(\"Done.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Create Merged Molecules"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will start with ligand 1 in the binding site. Create merged molecules between ligands 1 and 2, and 1 and 3:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the mapping and the inverse mapping\n",
    "mapping_3_1 = BSS.Align.matchAtoms(ligand_1_aligned, ligand_3)\n",
    "inv_mapping = {v: k for k, v in mapping_3_1.items()}\n",
    "\n",
    "# Align\n",
    "ligand_3_1_aligned = BSS.Align.rmsdAlign(ligand_3, ligand_1_aligned, inv_mapping)\n",
    "\n",
    "# Merge the two ligands based on the mapping.\n",
    "merged_3_1_aligned = BSS.Align.merge(ligand_1_aligned, ligand_3_1_aligned, mapping_3_1)\n",
    "\n",
    "#Create complex and free merged system\n",
    "\n",
    "complx_3_1 = merged_3_1_aligned + rna_bases \n",
    "\n",
    "merged_3_1 = merged_3_1_aligned \n",
    "\n",
    "#Solvate\n",
    "\n",
    "complx_3_1_sol = BSS.Solvent.tip3p(molecule= complx_3_1, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "merged_3_1_sol = BSS.Solvent.tip3p(molecule= merged_3_1, box=3 * [box_length], ion_conc = 0.15)\n",
    "# create the mapping and the inverse mapping\n",
    "mapping_2_1 = BSS.Align.matchAtoms(ligand_1_aligned, ligand_2)\n",
    "inv_mapping = {v: k for k, v in mapping_2_1.items()}\n",
    "\n",
    "# Align\n",
    "ligand_2_1_aligned = BSS.Align.rmsdAlign(ligand_2, ligand_1_aligned, inv_mapping)\n",
    "\n",
    "# Merge the two ligands based on the mapping.\n",
    "merged_2_1_aligned = BSS.Align.merge(ligand_1_aligned, ligand_2_1_aligned, mapping_2_1)\n",
    "\n",
    "#Create complex and free merged system\n",
    "\n",
    "complx_2_1 = merged_2_1_aligned + rna_bases\n",
    "\n",
    "merged_2_1 = merged_2_1_aligned\n",
    "\n",
    "#solvate\n",
    "\n",
    "complx_2_1_sol = BSS.Solvent.tip3p(molecule= complx_2_1, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "merged_2_1_sol = BSS.Solvent.tip3p(molecule= merged_2_1, box=3 * [box_length], ion_conc = 0.15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It is helpful to view the merged molecules to see that they are in the binding site and have hence aligned to the correct ligand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rna_test = complx_2_1\n",
    "view = BSS.Notebook.View(rna_test)\n",
    "view.system()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create merged moleucles with Ligand 2 in the binding site. Create merged molecules between ligands 2 and 3, and 2 and 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the mapping and the inverse mapping\n",
    "mapping_3_2 = BSS.Align.matchAtoms(ligand_2, ligand_3)\n",
    "inv_mapping = {v: k for k, v in mapping_3_2.items()}\n",
    "\n",
    "# Align\n",
    "ligand_3_2_aligned = BSS.Align.rmsdAlign(ligand_3, ligand_2, inv_mapping)\n",
    "\n",
    "# Merge the two ligands based on the mapping.\n",
    "merged_3_2_aligned = BSS.Align.merge(ligand_2, ligand_3_2_aligned, mapping_3_2)\n",
    "\n",
    "#Create complex and free merged system\n",
    "\n",
    "complx_3_2 = merged_3_2_aligned + rna_bases_2 \n",
    "\n",
    "merged_3_2 = merged_3_2_aligned \n",
    "\n",
    "#Solvate\n",
    "\n",
    "complx_3_2_sol = BSS.Solvent.tip3p(molecule= complx_3_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "merged_3_2_sol = BSS.Solvent.tip3p(molecule= merged_3_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "# create the mapping and the inverse mapping\n",
    "mapping_1_2 = BSS.Align.matchAtoms(ligand_2, ligand_1)\n",
    "inv_mapping = {v: k for k, v in mapping_1_2.items()}\n",
    "\n",
    "# Align\n",
    "ligand_1_2_aligned = BSS.Align.rmsdAlign(ligand_1, ligand_2, inv_mapping)\n",
    "\n",
    "# Merge the two ligands based on the mapping.\n",
    "merged_1_2_aligned = BSS.Align.merge(ligand_2, ligand_1_2_aligned, mapping_1_2)\n",
    "\n",
    "#Create complex and free merged system\n",
    "\n",
    "complx_1_2 = merged_1_2_aligned + rna_bases_2 \n",
    "\n",
    "merged_1_2 = merged_1_2_aligned \n",
    "\n",
    "#Solvate\n",
    "\n",
    "complx_1_2_sol = BSS.Solvent.tip3p(molecule= complx_1_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "merged_1_2_sol = BSS.Solvent.tip3p(molecule= merged_1_2, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we can create merged moleucles with Ligand 3 in the binding site. Create merged molecules between ligands 3 and 2, and 3 and 1:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the mapping and the inverse mapping\n",
    "mapping_2_3 = BSS.Align.matchAtoms(ligand_2, ligand_3)\n",
    "inv_mapping_2_3 = {v: k for k, v in mapping_2_3.items()}\n",
    "\n",
    "# Align\n",
    "ligand_2_3_aligned = BSS.Align.rmsdAlign(ligand_3, ligand_2, inv_mapping_2_3)\n",
    "\n",
    "# Merge the two ligands based on the mapping.\n",
    "merged_2_3_aligned = BSS.Align.merge(ligand_2, ligand_2_3_aligned, mapping_2_3)\n",
    "\n",
    "#Create complex and free merged system\n",
    "\n",
    "complx_2_3 = merged_2_3_aligned + rna_bases_3 \n",
    "\n",
    "merged_2_3 = merged_2_3_aligned \n",
    "\n",
    "#solvate\n",
    "\n",
    "complx_2_3_sol = BSS.Solvent.tip3p(molecule= complx_2_3, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "merged_2_3_sol = BSS.Solvent.tip3p(molecule= merged_2_3, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "# create the mapping and the inverse mapping\n",
    "mapping_1_3 = BSS.Align.matchAtoms(ligand_1, ligand_3)\n",
    "inv_mapping_1_3 = {v: k for k, v in mapping_1_3.items()}\n",
    "\n",
    "# Align\n",
    "ligand_1_3_aligned = BSS.Align.rmsdAlign(ligand_3, ligand_1, inv_mapping_1_3)\n",
    "\n",
    "# Merge the two ligands based on the mapping.\n",
    "merged_1_3_aligned = BSS.Align.merge(ligand_1, ligand_1_3_aligned, mapping_1_3)\n",
    "\n",
    "#Create complex and free merged system\n",
    "\n",
    "complx_1_3 = merged_1_3_aligned + rna_bases_3 \n",
    "\n",
    "merged_1_3 = merged_1_3_aligned \n",
    "\n",
    "#solvate\n",
    "\n",
    "complx_1_3_sol = BSS.Solvent.tip3p(molecule= complx_1_3, box=3 * [box_length], ion_conc = 0.15)\n",
    "\n",
    "merged_1_3_sol = BSS.Solvent.tip3p(molecule= merged_1_3, box=3 * [box_length], ion_conc = 0.15)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generate .pert files ready for RBFE calculations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have merged molecules for Ligand 1:Ligand 2 and Ligand 1:Ligand 3 we can generate the input files (.pert) for the RBFE calculations using BSS.FreeEnergy.Relative.\n",
    "\n",
    "- Note: in the extra options directory I have specified nmoves = 25000, ncycles = 50, timestep = 4 fs and hydrogen mass repartitioning factor = 3.0. This results in 5 ns simulation time and is done to ensure full engagement of GPUs.\n",
    "- Currently the lamba window configuration files require 'femtosecond' rather than 'fs'. NO FIX CURRENTLY"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def pert_gen(ligand, system, name)\n",
    "    ### Settings.\n",
    "    minim_steps = 250\n",
    "    runtime_short_nvt = 5 # ps\n",
    "    runtime_nvt = 50 # ps RESET to 50 after TESTING ! \n",
    "    runtime_npt = 200 # ps RESET to 200 after TESTING !\n",
    "\n",
    "    ### preamble. tmp_dir should at some point be derived using os.environ.get(\"\")\n",
    "\n",
    "    tmp_dir = \"temp\"\n",
    "\n",
    "    amber_home = \"/usr/local/amber22\"\n",
    "    pmemd_path = amber_home + \"/bin/pmemd.cuda\" \n",
    "\n",
    "    #################\n",
    "    lig_p_solvated = ligand\n",
    "    system_solvated = system\n",
    "\n",
    "    print (\"lig_p_solvated is =\", lig_p_solvated)\n",
    "    print (\"original is =\", ligand)\n",
    "\n",
    "    freenrg_protocol = BSS.Protocol.FreeEnergy(\n",
    "        num_lam=16, runtime= 5 * BSS.Units.Time.nanosecond\n",
    "    )\n",
    "    # set up a bound folder with standard settings.\n",
    "    # Use solvation to prep the bound leg\n",
    "    print(\"Bound..\")\n",
    "    workdir = f\"prep_bases/rbfe_outputs/{name}\"\n",
    "    BSS.FreeEnergy.Relative(\n",
    "        system_solvated,\n",
    "        freenrg_protocol,\n",
    "        engine=f\"SOMD\",\n",
    "        work_dir=workdir + \"/bound\",\n",
    "        extra_options={ \"nmoves\" : 25000, \"ncycles\" : 50, \"timestep\" : 4 * BSS.Units.Time.femtosecond, \"hydrogen mass repartitioning factor\" : 3.0},\n",
    "    )\n",
    "\n",
    "    # set up a free folder.\n",
    "    print(\"Free..\")\n",
    "    BSS.FreeEnergy.Relative(\n",
    "        lig_p_solvated,\n",
    "        freenrg_protocol,\n",
    "        engine=f\"SOMD\",\n",
    "        work_dir=workdir + \"/free\",\n",
    "        extra_options={ \"nmoves\" : 25000, \"ncycles\" : 50, \"timestep\" : 4 * BSS.Units.Time.femtosecond, \"hydrogen mass repartitioning factor\" : 3.0},\n",
    "    )\n",
    "pert_gen(merged_2_1_sol, complx_2_1_sol, \"lig_2_1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4b. Analyse RBFE"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Calculate the RBFE and compare to a literature value."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pmf_free, overlap_matrix_free = BSS.FreeEnergy.Relative.analyse(f'prep_bases/rbfe_outputs/lig_1_2_r1/free')\n",
    "pmf_bound, overlap_matrix_bound = BSS.FreeEnergy.Relative.analyse(f'prep_bases/rbfe_outputs/lig_1_2_r1/bound')\n",
    "freenrg_rel = BSS.FreeEnergy.Relative.difference(pmf_bound, pmf_free)\n",
    "\n",
    "print(f\"pmf_free is: {pmf_free[-1][1]} and the MBAR statistical uncertainty is {pmf_free[-1][2]} .\")\n",
    "print(f\"pmf_bound is: {pmf_bound[-1][1]} and the MBAR statistical uncertainty is {pmf_bound[-1][2]} .\")\n",
    "print(f\"The RBFE is {freenrg_rel[0]} and the error is {freenrg_rel[1]}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Below is a function to calculate the RBFE between two literature ligands and convert the value to kcal/mol."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "def kd_to_deltaG(kd, temperature):\n",
    "    # Gas constant in J/(mol·K)\n",
    "    R = 8.314\n",
    "\n",
    "    # Convert temperature to Kelvin if it's given in Celsius\n",
    "    if temperature < 273.15:\n",
    "        temperature += 273.15\n",
    "\n",
    "    # Calculate delta G in joules/mol\n",
    "    delta_G_joules = -R * temperature * math.log(kd)\n",
    "\n",
    "    # Convert delta G to kilocalories/mol\n",
    "    delta_G_kcal = delta_G_joules * 0.000239\n",
    "\n",
    "    return delta_G_kcal\n",
    "\n",
    "# Example usage\n",
    "kd_value = 0.5 * 1e-6  # Kd in micro M\n",
    "temperature_value = 300  # Temperature in Kelvin\n",
    "\n",
    "deltaG_result = kd_to_deltaG(kd_value, temperature_value)\n",
    "print(f\"Free Energy (ΔG) of 1 = {deltaG_result} kcal/mol\")\n",
    "\n",
    "kd_value_2 = 0.4 * 1e-6  # Kd in micro M\n",
    "temperature_value_2 = 300  # Temperature in Kelvin\n",
    "\n",
    "deltaG_result_2 = kd_to_deltaG(kd_value_2, temperature_value_2)\n",
    "print(f\"Free Energy (ΔG) of 2 = {deltaG_result_2} kcal/mol\")\n",
    "\n",
    "diff = deltaG_result_2 - deltaG_result\n",
    "print (f\"diff = {diff} kcal/mol\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "bsstutorials",
   "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"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}