{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# BespokeFit: Using bespoke force field parameters with OpenFE protocols\n",
    "\n",
    "This tutorial gives a step-by-step guide on the use of OpenFF-BespokeFit generated force field parameters with OpenFE protocols. Here we will focus on relative\n",
    "binding free energy (RBFE) calculations, but this strategy can be used with any OpenMM-based OpenFE protocol. \n",
    "\n",
    "Here we will assume you have succesfully planned your RBFE campaign using OpenFE by following the [showcase](http://try.openfree.energy/) or [RBFE tutorial](https://docs.openfree.energy/en/stable/tutorials/rbfe_cli_tutorial.html#rbfe-cli-tutorial) and will be using the `TYK2` test system in this example with the planned network provided for you at `assets/ligand_network.graphml`. You should have also already generated a set of bespoke parameters for your ligand series following the [BespokeFit production guide](https://docs.openforcefield.org/projects/bespokefit/en/latest/getting-started/quick-start.html#production-fits) and combined the parameters into a single SMIRNOFF style `offxml` file following the [gathering results guide](https://docs.openforcefield.org/projects/bespokefit/en/latest/users/bespoke-results.html) an example bespoke force field is provied for you at `assets/bespoke-1.3.0-default.offxml`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reading input files\n",
    "\n",
    "Before creating the OpenFE simulation inputs we first we need to load up our planned ligand network of transformations and our bespoke force field file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from gufe import LigandNetwork\n",
    "from openff.toolkit import ForceField\n",
    "\n",
    "# load our network file\n",
    "ligand_network = LigandNetwork.from_graphml(\n",
    "    open(\"assets/ligand_network.graphml\", \"r\").read()\n",
    ")\n",
    "# load the bespoke force field\n",
    "bespoke_force_field = ForceField(\"assets/bespoke-1.3.0-default.offxml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running the simulations\n",
    "\n",
    "We are now ready to build a new set of transformations which can be executed locally using the [OpenFE quickrun CLI](https://docs.openfree.energy/en/stable/tutorials/rbfe_cli_tutorial.html#running-the-simulations). We will be following the [python API tutorial](https://docs.openfree.energy/en/stable/tutorials/rbfe_python_tutorial.html#Creating-a-Protocol) to create the transformations, you should check that documentation for more details on the next steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create the OpenFE RBFE protocol using our bespoke force field\n",
    "from openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol\n",
    "import openfe\n",
    "\n",
    "\n",
    "# create the default protocol settings\n",
    "settings = RelativeHybridTopologyProtocol.default_settings()\n",
    "# add our new force field as a string\n",
    "# this avoids the need to move the file around when executing the transformations\n",
    "settings.forcefield_settings.small_molecule_forcefield = bespoke_force_field.to_string()\n",
    "\n",
    "# create the protocol\n",
    "protocol = RelativeHybridTopologyProtocol(settings)\n",
    "\n",
    "# create the solvent and protein components\n",
    "solvent = openfe.SolventComponent()\n",
    "protein = openfe.ProteinComponent.from_pdb_file(\"assets/tyk2_protein.pdb\")\n",
    "\n",
    "# follow the tutorial to create the AlchemicalNetwork\n",
    "transformations = []\n",
    "for mapping in ligand_network.edges:\n",
    "    for leg in ['solvent', 'complex']:\n",
    "        # use the solvent and protein created above\n",
    "        sysA_dict = {'ligand': mapping.componentA,\n",
    "                     'solvent': solvent}\n",
    "        sysB_dict = {'ligand': mapping.componentB,\n",
    "                     'solvent': solvent}\n",
    "\n",
    "        if leg == 'complex':\n",
    "            sysA_dict['protein'] = protein\n",
    "            sysB_dict['protein'] = protein\n",
    "\n",
    "        # we don't have to name objects, but it can make things (like filenames) more convenient\n",
    "        sysA = openfe.ChemicalSystem(sysA_dict, name=f\"{mapping.componentA.name}_{leg}\")\n",
    "        sysB = openfe.ChemicalSystem(sysB_dict, name=f\"{mapping.componentB.name}_{leg}\")\n",
    "\n",
    "        prefix = \"rbfe_\"  # prefix is only to exactly reproduce CLI\n",
    "\n",
    "        transformation = openfe.Transformation(\n",
    "            stateA=sysA,\n",
    "            stateB=sysB,\n",
    "            mapping={'ligand': mapping},\n",
    "            protocol=protocol,  # use protocol created above\n",
    "            name=f\"{prefix}{sysA.name}_{sysB.name}\"\n",
    "        )\n",
    "        transformations.append(transformation)\n",
    "\n",
    "network = openfe.AlchemicalNetwork(transformations)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can now write out each of the transformations to disk for independent execution: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "# first we create the directory\n",
    "transformation_dir = pathlib.Path(\"transformations\")\n",
    "transformation_dir.mkdir(exist_ok=True)\n",
    "\n",
    "# then we write out each transformation\n",
    "for transformation in network.edges:\n",
    "    transformation.dump(transformation_dir / f\"{transformation.name}.json\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Recap\n",
    "\n",
    "So to recap the workflow can be reduced to the following steps:\n\n",
    "- Plan the RBFE network\n",
    "- Create a single SMIRNOFF style force field with all of the bespoke parameters for the network using the BespokeFit `combine` CLI\n",
    "- Store the force field as a string in the OpenFE protocol under the `settings.forcefield_settings.small_molecule_forcefield` field\n",
    "- Use these settings to create the protcol and create the AlchemicalNetwork following the normal steps\n",
    "\n",
    "Hopefully its clear that this strategy can be applied to any bespoke parameters you wish to add to the force field not just those from BespokeFit, simply edit your base SMIRNOFF style force field using the [OpenFF-Toolkit](https://docs.openforcefield.org/en/latest/examples/openforcefield/openff-toolkit/forcefield_modification/forcefield_modification.html#modifying-a-smirnoff-force-field), set it in the protocol and simulate! \n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "asapdiscovery",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}