{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Methods for Topology solvation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "This example explores some different ways to solvate an OpenFF Topology. See also the \"Generating and Parametrizing multi-component systems\", \"Solvating and equilibrating a ligand in a box of water\", and the Toolkit Showcase."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import Optional\n",
    "\n",
    "import nglview\n",
    "import numpy\n",
    "import openmm\n",
    "import openmm.app\n",
    "import openmm.unit\n",
    "from openff.toolkit import ForceField, Molecule, Topology\n",
    "from openff.units import Quantity, unit\n",
    "from openff.units.openmm import ensure_quantity\n",
    "from openmm.app import Topology as OpenMMTopology\n",
    "\n",
    "from openff.interchange import Interchange"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "## Methods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "OPENMM_IONS = {\n",
    "    \"Li+\": \"[#3+1]\",\n",
    "    \"Na+\": \"[#11+1]\",\n",
    "    \"K+\": \"[#19+1]\",\n",
    "    \"Rb+\": \"[#37+1]\",\n",
    "    \"Cs+\": \"[#55+1]\",\n",
    "    \"F-\": \"[#9-1]\",\n",
    "    \"Cl-\": \"[#17-1]\",\n",
    "    \"Br-\": \"[#35-1]\",\n",
    "    \"I-\": \"[#53-1]\",\n",
    "}\n",
    "\n",
    "\n",
    "def visualize_all(interchange: Interchange) -> nglview.NGLWidget:\n",
    "    view = interchange.visualize()\n",
    "    view.clear_representations()\n",
    "    view.add_representation(\"ball+stick\", selection=\"all\")\n",
    "    return view\n",
    "\n",
    "\n",
    "def solvate_topology(\n",
    "    topology: Topology,\n",
    "    method: str = \"pdbfixer\",\n",
    "    box_vectors: Optional[Quantity] = Quantity(5.0 * numpy.ones(3), unit.nanometer),\n",
    "    **kwargs,\n",
    ") -> Topology:\n",
    "    if method in [\"pdbfixer\", \"openmm\"]:\n",
    "        boxSize = openmm.unit.Quantity(\n",
    "            openmm.Vec3(*box_vectors.m_as(unit.nanometer)), openmm.unit.nanometer\n",
    "        )\n",
    "\n",
    "        if method == \"pdbfixer\":\n",
    "            openmm_topology, openmm_positions = _solvate_pdbfixer(\n",
    "                topology.to_openmm(),\n",
    "                topology.get_positions().to_openmm(),\n",
    "                boxSize=boxSize,\n",
    "                **kwargs,\n",
    "            )\n",
    "        else:\n",
    "            openmm_topology, openmm_positions = _solvate_openmm(\n",
    "                topology.to_openmm(),\n",
    "                topology.get_positions().to_openmm(),\n",
    "                boxSize=boxSize,\n",
    "                **kwargs,\n",
    "            )\n",
    "\n",
    "        unique_molecules: List[Molecule] = [*topology.unique_molecules]\n",
    "        unique_molecules.append(Molecule.from_mapped_smiles(\"[H:2][O:1][H:3]\"))\n",
    "\n",
    "        if \"positiveIon\" in kwargs:\n",
    "            unique_molecules.append(\n",
    "                Molecule.from_smiles(OPENMM_IONS[kwargs[\"positiveIon\"]])\n",
    "            )\n",
    "\n",
    "        if \"negativeIon\" in kwargs:\n",
    "            unique_molecules.append(\n",
    "                Molecule.from_smiles(OPENMM_IONS[kwargs[\"negativeIon\"]])\n",
    "            )\n",
    "\n",
    "        new_topology = Topology.from_openmm(\n",
    "            openmm_topology,\n",
    "            unique_molecules=unique_molecules,\n",
    "        )\n",
    "\n",
    "        new_topology.set_positions(ensure_quantity(openmm_positions, \"openff\"))\n",
    "\n",
    "        return new_topology\n",
    "\n",
    "\n",
    "def _solvate_pdbfixer(\n",
    "    topology: OpenMMTopology,\n",
    "    positions: openmm.unit.Quantity,\n",
    "    **kwargs,\n",
    ") -> tuple[OpenMMTopology, openmm.unit.Quantity]:\n",
    "    \"\"\"\n",
    "    Add solvent and ions using PDBFixer.\n",
    "\n",
    "    https://htmlpreview.github.io/?https://github.com/openmm/pdbfixer/blob/master/Manual.html\n",
    "\n",
    "    \"\"\"\n",
    "    import pdbfixer\n",
    "\n",
    "    with open(\"_tmp.pdb\", \"w\") as _file:\n",
    "        openmm.app.PDBFile.writeFile(topology, positions, _file)\n",
    "\n",
    "    pdb_object = pdbfixer.PDBFixer(\"_tmp.pdb\")\n",
    "    pdb_object.addSolvent(**kwargs)\n",
    "\n",
    "    return pdb_object.topology, pdb_object.positions\n",
    "\n",
    "\n",
    "def _solvate_openmm(\n",
    "    topology: OpenMMTopology,\n",
    "    positions: openmm.unit.Quantity,\n",
    "    box_vectors: openmm.unit.Quantity,\n",
    "    forcefield: Optional[openmm.app.ForceField] = None,\n",
    "    **kwargs,\n",
    ") -> tuple[OpenMMTopology, openmm.unit.Quantity]:\n",
    "    if not forcefield:\n",
    "        import pdbfixer\n",
    "\n",
    "        forcefield = pdbfixer.PDBFixer._createForceField(topology)\n",
    "\n",
    "    modeller = openmm.app.Modeller(topology, positions)\n",
    "    modeller.addSolvent(\n",
    "        forcefield,\n",
    "        **kwargs,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {
    "tags": []
   },
   "source": [
    "## Solvating and Parametrizing"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "Load the force field:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "sage_ff14sb = ForceField(\"openff-2.0.0.offxml\", \"ff14sb_off_impropers_0.0.3.offxml\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "Load the solute from a PDB file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "peptide = Molecule.from_polymer_pdb(\"ace-a5ca5-nme.pdb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "Solvate:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "solvated_topology = solvate_topology(\n",
    "    peptide.to_topology(),\n",
    "    box_vectors=Quantity(5.0 * numpy.ones(3), unit.nanometer),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "12",
   "metadata": {},
   "source": [
    "Parametrize and visualize:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "interchange = Interchange.from_smirnoff(sage_ff14sb, solvated_topology)\n",
    "visualize_all(interchange)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14",
   "metadata": {},
   "source": [
    "And finally, some alternative solvation strategies:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "15",
   "metadata": {},
   "outputs": [],
   "source": [
    "solvated_topology = solvate_topology(\n",
    "    peptide.to_topology(),\n",
    "    box_vectors=Quantity([10, 4, 4], unit.nanometer),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16",
   "metadata": {},
   "outputs": [],
   "source": [
    "solvated_topology = solvate_topology(\n",
    "    peptide.to_topology(),\n",
    "    box_vectors=Quantity(5.0 * numpy.ones(3), unit.nanometer),\n",
    "    positiveIon=\"K+\",\n",
    "    negativeIon=\"Br-\",\n",
    "    ionicStrength=2.0 * openmm.unit.molar,\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.4"
  },
  "vscode": {
   "interpreter": {
    "hash": "86c9b142c8dc60dd36d17e2a57efabbd2ed015b9d3db80dd77f3e0894d5aea85"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}