{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b89fdb9a-24fe-4aeb-9390-a359dae6cc9b",
   "metadata": {},
   "source": [
    "# Create an Alchemical Network"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b11a786-0fab-4990-a2aa-242067f89287",
   "metadata": {},
   "source": [
    "The final setup step is to compile all the different bits of information into a description of a single simulation campaign. This description takes the form of an [Alchemical Network]. Similarly to `LigandNetwork`, the `AlchemicalNetwork` class is a graph of all the transformations in the campaign; however, in an `AlchemicalNetwork`, these transformations include all the information needed to perform the transformation. By contrast, a `LigandNetwork` includes only the ligands themselves.\n",
    "\n",
    "[Alchemical Network]: https://docs.openfree.energy/en/stable/reference/api/generated/openfe.AlchemicalNetwork.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "471531a1-a351-413b-b147-794f9af4e0f8",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ba6d0d2f-03df-463f-8976-c15a2bdf65e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import openfe, rdkit.Chem\n",
    "from openff.units import unit\n",
    "from openfe.protocols.openmm_rfe import RelativeHybridTopologyProtocol"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bfee4c7-0aaf-49a6-b35a-bdbe99444bab",
   "metadata": {},
   "source": [
    "This cookbook assumes you've already loaded a `LigandNetwork`. For more information, see [Generate a Ligand Network Automatically]:\n",
    "\n",
    "[Generate a Ligand Network Automatically]: https://docs.openfree.energy/en/stable/cookbook/generate_ligand_network.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4512c4d6-720f-4008-af53-a0b56b2ee1e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "ligand_network = openfe.ligand_network_planning.generate_minimal_spanning_network(\n",
    "    ligands=[\n",
    "        openfe.SmallMoleculeComponent(mol) \n",
    "        for mol in rdkit.Chem.SDMolSupplier(\n",
    "            \"assets/somebenzenes.sdf\", \n",
    "            removeHs=False,\n",
    "        )\n",
    "    ],\n",
    "    mappers=[openfe.setup.LomapAtomMapper()],\n",
    "    scorer=openfe.lomap_scorers.default_lomap_score,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "405fe819-b4a6-408b-8623-041492c84dfb",
   "metadata": {},
   "source": [
    "This cookbook assumes you've already loaded a `Protocol`. For more information, see [Choose and Configure a Protocol]:\n",
    "\n",
    "[Choose and Configure a Protocol]: https://docs.openfree.energy/en/stable/cookbook/choose_protocol.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "940e2e19-4e7d-4853-ab5e-54626b3384d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "protocol = RelativeHybridTopologyProtocol(RelativeHybridTopologyProtocol.default_settings())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0908d2b2-421d-4b51-8b35-d2cca5e25a95",
   "metadata": {},
   "source": [
    "This cookbook assumes you've already loaded a solvent and all the other chemical components of the system, including any proteins or cofactors. For more information, see [Loading your data into Components]:\n",
    "\n",
    "[Loading your data into Components]: https://docs.openfree.energy/en/stable/cookbook/loading_molecules.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d3054cfc-5040-4882-b3a3-92aba7ab5d1d",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "solvent = openfe.SolventComponent(\n",
    "    ion_concentration=0.15 * unit.molar\n",
    ")\n",
    "protein = openfe.ProteinComponent.from_pdb_file(\n",
    "    \"assets/t4_lysozyme.pdb\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9249e15a-f4cb-425d-8fde-f28c79de670e",
   "metadata": {},
   "source": [
    "## Create the Alchemical Network"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b2f6271-3c89-4cd6-8283-e617d535b910",
   "metadata": {},
   "source": [
    "### Automatically"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc13516d-92b6-44ac-a3eb-20e7f9e00631",
   "metadata": {},
   "source": [
    "The `LigandNetwork.to_rbfe_alchemical_network()` method makes constructing alchemical networks for relative binding free energy calculations very simple: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "733dee64-d23a-4eb0-87cf-c03065b8d2a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "alchemical_network_auto = ligand_network.to_rbfe_alchemical_network(\n",
    "    solvent=solvent,\n",
    "    protein=protein,\n",
    "    protocol=protocol,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35774b4f-83e7-4235-960b-aa591e20e016",
   "metadata": {},
   "source": [
    "### Manually"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0da5558b-3894-4c11-aafd-e4478972e161",
   "metadata": {},
   "source": [
    "If your needs are not catered to by the above method, you can instead loop over the `LigandNetwork` edges and manually create the `Transformation` objects for each of them. This gives you full control over the entire network. For more information, see [Under the Hood]:\n",
    "\n",
    "[Under the Hood]: https://docs.openfree.energy/en/stable/cookbook/under_the_hood.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c49f36b1-c3c1-425e-b1de-24da25ed01c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In an RBFE, each edge includes two \"legs\": \n",
    "# one for the ligand complexed to the protein, \n",
    "# and the other for the ligand free in solution\n",
    "legs = {\n",
    "    \"solvent\": {\n",
    "        # Specify the components common to all systems in this leg\n",
    "        'solvent': solvent,\n",
    "    },\n",
    "    \"complex\": {\n",
    "        # Specify the components common to all systems in this leg\n",
    "        'solvent': solvent,\n",
    "        'protein': protein,\n",
    "    }\n",
    "}\n",
    "transformations = []\n",
    "\n",
    "for mapping in ligand_network.edges:\n",
    "    for leg, common_components in legs.items():\n",
    "        system_a = openfe.ChemicalSystem(\n",
    "            {\n",
    "                'ligand': mapping.componentA,\n",
    "                **common_components,\n",
    "            },\n",
    "            name=f\"{mapping.componentA.name}_{leg}\"\n",
    "        )\n",
    "        \n",
    "        system_b = openfe.ChemicalSystem(\n",
    "            {\n",
    "                'ligand': mapping.componentB,\n",
    "                **common_components,\n",
    "            },\n",
    "            name=f\"{mapping.componentB.name}_{leg}\"\n",
    "        )\n",
    "    \n",
    "        transformation = openfe.Transformation(\n",
    "            stateA=system_a,\n",
    "            stateB=system_b,\n",
    "            mapping=mapping,\n",
    "            protocol=protocol,\n",
    "            # Using the same name as to_rbfe_alchemical_network()\n",
    "            name=f\"easy_rbfe_{system_a.name}_{system_b.name}\"\n",
    "        )\n",
    "    \n",
    "        transformations.append(transformation)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9772b2b4-7194-456e-9288-fbe82937f187",
   "metadata": {},
   "source": [
    "Finally, combine the transformations into an Alchemical Network:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "fc6cac68-bdcd-4d51-b4e1-f9bb7169c6f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "alchemical_network = openfe.AlchemicalNetwork(transformations)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b623d82e-1069-4c69-81af-ccc795ba5690",
   "metadata": {},
   "source": [
    "We can confirm that this produces identical results to the previous strategy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c9baa54c-513a-486a-a06a-fd8fd5d19fa5",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "assert alchemical_network == ligand_network.to_rbfe_alchemical_network(\n",
    "    solvent=solvent,\n",
    "    protein=protein,\n",
    "    protocol=protocol,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b17669f6-a205-420d-96c3-1a9c92d0e13b",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "## Write the Alchemical Network to Disk"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "682b45eb-4c79-4a68-b0a8-534544c0f0a7",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "source": [
    "While the `AlchemicalNetwork` class itself has no on-disk representation, the `Transformation` instances that compose it do. Write an alchemical network to disk by iterating over its edges. For more information, see [Dumping a Transformation to JSON]:\n",
    "\n",
    "[Dumping a Transformation to JSON]: https://docs.openfree.energy/en/stable/cookbook/dumping_transformation.html"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2da4dafb-20d6-4ce7-ab83-4690528097d6",
   "metadata": {
    "slideshow": {
     "slide_type": ""
    },
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "transformations_dir = Path(\"transformations\")\n",
    "transformations_dir.mkdir(exist_ok=True)\n",
    "\n",
    "for n, transformation in enumerate(alchemical_network.edges):\n",
    "    transformation_name = transformation.name or transformation.key\n",
    "    transformation.dump(transformations_dir / f\"{n}_{transformation_name}.json\")"
   ]
  }
 ],
 "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
}