{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b89fdb9a-24fe-4aeb-9390-a359dae6cc9b",
   "metadata": {},
   "source": [
    "# Relative Free Energy Alchemical Network Planners"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b11a786-0fab-4990-a2aa-242067f89287",
   "metadata": {},
   "source": [
    "To simplify the process of creating an [Alchemical Network] for Relative Free Energy calculations, openfe offers two convenience `Alchemical Network Planner` classes; [RBFEAlchemicalNetworkPlanner] and [RHFEAlchemicalNetworkPlanner].\n",
    "\n",
    "[Alchemical Network]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.AlchemicalNetwork.html\n",
    "[RBFEAlchemicalNetworkPlanner]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.setup.RBFEAlchemicalNetworkPlanner.html#openfe.setup.RBFEAlchemicalNetworkPlanner\n",
    "[RHFEAlchemicalNetworkPlanner]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.setup.RHFEAlchemicalNetworkPlanner.html#openfe.setup.RHFEAlchemicalNetworkPlanner "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d449f31",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "To create an `AlchemicalNetwork` using the Alchemical Planners, the following steps must be followed:\n",
    "\n",
    "1. Create appropriate `ChemicalComponents`\n",
    "2. Choosing an appropriate atom mapper, scorer, and network generator\n",
    "3. Create an `Alchemical Planner`\n",
    "4. Create the `AlchemicalNetwork`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "471531a1-a351-413b-b147-794f9af4e0f8",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Relative Binding Free Energy (RBFE)\n",
    "\n",
    "Here we demonstrate how to use the [RBFEAlchemicalNetworkPlanner] to create an RBFE [AlchemicalNetwork]\n",
    "\n",
    "[AlchemicalNetwork]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.AlchemicalNetwork.html\n",
    "[RBFEAlchemicalNetworkPlanner]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.setup.RBFEAlchemicalNetworkPlanner.html#openfe.setup.RBFEAlchemicalNetworkPlanner"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1b9f171b",
   "metadata": {},
   "source": [
    "### Creating `ChemicalComponents`\n",
    "\n",
    "First we define the `ChemicalComponents` which we want to represent in our `AlchemicalNetwork`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ba6d0d2f-03df-463f-8976-c15a2bdf65e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from rdkit import Chem\n",
    "from openfe import SmallMoleculeComponent, SolventComponent, ProteinComponent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "ceeff0fb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Some SmallMoleculeComponents defining the ligands we want to transform between\n",
    "ligands = [\n",
    "    SmallMoleculeComponent(m)\n",
    "    for m in Chem.SDMolSupplier(\"assets/somebenzenes.sdf\", removeHs=False)\n",
    "]\n",
    "\n",
    "# The ProteinComponent which will be present in our complex simulations\n",
    "protein = ProteinComponent.from_pdb_file(\"assets/t4_lysozyme.pdb\")\n",
    "\n",
    "# The SolventComponent defining our system's solvation\n",
    "solvent = SolventComponent()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bfee4c7-0aaf-49a6-b35a-bdbe99444bab",
   "metadata": {},
   "source": [
    "### Choosing an atom mapper, scorer and network generator\n",
    "\n",
    "Here we choose to use the [KartografAtomMapper] alongside the [Lomap scorer].\n",
    "\n",
    "For our network we choose to create a minimum spanning network using [generate_minimal_spanning_network].\n",
    "\n",
    "[KartografAtomMapper]: https://kartograf.readthedocs.io/en/latest/api/kartograf.mappers.html\n",
    "[Lomap scorer]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.setup.atom_mapping.lomap_scorers.default_lomap_score.html\n",
    "[generate_minimal_spanning_network]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.setup.ligand_network_planning.generate_minimal_spanning_network.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "1284865e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openfe.setup.atom_mapping import KartografAtomMapper\n",
    "from openfe.setup.atom_mapping.lomap_scorers import default_lomap_score\n",
    "from openfe.setup.ligand_network_planning import generate_minimal_spanning_network"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e96f95fb",
   "metadata": {},
   "source": [
    "### Building the RBFEAlchemicalNetworkPlanner and getting an AlchemicalNetwork\n",
    "\n",
    "Finally we can use the `RBFEAlchemicalNetworkPlanner` to create an `AlchemicalNetwork`.\n",
    "\n",
    "Internally the Network Planner will:\n",
    "\n",
    "1. Create a `LigandNetwork` (accessible under the `._ligand_network` attribute).\n",
    "2. Create two `Transformation` per edge of the `LigandNetwork`, one for each of the complex and solvent legs of the RBFE alchemical cycle.\n",
    "3. Add all `Transformation`s to a single `AlchemicalNetwork`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "4512c4d6-720f-4008-af53-a0b56b2ee1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from openfe.setup.alchemical_network_planner import RBFEAlchemicalNetworkPlanner\n",
    "\n",
    "# Build the planner\n",
    "alchem_planner = RBFEAlchemicalNetworkPlanner(\n",
    "    name=\"project\", \n",
    "    mappers=[KartografAtomMapper()],\n",
    "    mapping_scorer=default_lomap_score,\n",
    "    ligand_network_planner=generate_minimal_spanning_network,\n",
    ")\n",
    "\n",
    "# Create the AlchemicalNetwork\n",
    "alchemical_network = alchem_planner(\n",
    "    ligands=ligands,\n",
    "    solvent=solvent,\n",
    "    protein=protein\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "405fe819-b4a6-408b-8623-041492c84dfb",
   "metadata": {},
   "source": [
    "## Relative Hydration Free Energy (RHFE)\n",
    "\n",
    "Similarly, an RHFE `AlchemicalNetwork` can be created using the [RHFEAlchemicalNetworkPlanner].\n",
    "\n",
    "[RHFEAlchemicalNetworkPlanner]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.setup.RHFEAlchemicalNetworkPlanner.html#openfe.setup.RHFEAlchemicalNetworkPlanner "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc5f6bcb",
   "metadata": {},
   "source": [
    "### Building the RHFEAlchemicalNetworkPlanner and getting an AlchemicalNetwork\n",
    "\n",
    "Here we use the same ChemicalComponents, atom mapper, scorer, and network planner as used in the RBFE section above.\n",
    "\n",
    "The `RHFEAlchemicalNetworkPlanner` will:\n",
    "\n",
    "1. Create a new `LigandNetwork`.\n",
    "2. Create two `Transformation` per `LigandNetwork` edge, one for the solvent and vacuum legs of the transformation, setting appropriate nonbonded methods and cutoffs for the periodic and nonperiodic simulations.\n",
    "3. Combine all the `Transformation`s into an `AlchemicalNetwork`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d3054cfc-5040-4882-b3a3-92aba7ab5d1d",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from openfe.setup.alchemical_network_planner import RHFEAlchemicalNetworkPlanner\n",
    "\n",
    "# Build the planner\n",
    "alchem_planner = RHFEAlchemicalNetworkPlanner(\n",
    "    name=\"project\", \n",
    "    mappers=[KartografAtomMapper()],\n",
    "    mapping_scorer=default_lomap_score,\n",
    "    ligand_network_planner=generate_minimal_spanning_network,\n",
    ")\n",
    "\n",
    "# Create the AlchemicalNetwork\n",
    "alchemical_network = alchem_planner(\n",
    "    ligands=ligands,\n",
    "    solvent=solvent,\n",
    ")"
   ]
  }
 ],
 "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.11.6"
  },
  "widgets": {
   "application/vnd.jupyter.widget-state+json": {
    "state": {},
    "version_major": 2,
    "version_minor": 0
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}