{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Simulate an Interchange with OpenMM\n",
    "\n",
    "In this example, we'll quickly construct an `Interchange` and then run a simulation in OpenMM. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We need an `Interchange` to get started, so let's put that together quickly. For more explanation on this process, take a look at the [packed_box] and [protein_ligand] examples.\n",
    "\n",
    "[packed_box]: https://github.com/openforcefield/openff-interchange/tree/main/examples/packed_box\n",
    "[protein_ligand]: https://github.com/openforcefield/openff-interchange/tree/main/examples/protein_ligand"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "import mdtraj as md\n",
    "import nglview\n",
    "import openmm\n",
    "from openff.toolkit.topology import Molecule, Topology\n",
    "from openff.toolkit.typing.engines.smirnoff import ForceField\n",
    "from openff.toolkit.utils import get_data_file_path\n",
    "from pandas import read_csv\n",
    "\n",
    "# This prepared PDB file from the toolkit's test suite is a box of solvents\n",
    "pdb_path = get_data_file_path(\"systems/packmol_boxes/propane_methane_butanol_0.2_0.3_0.5.pdb\")\n",
    "molecules = [Molecule.from_smiles(smi) for smi in [\"CCC\", \"C\", \"CCCCO\"]]\n",
    "\n",
    "# The OpenFF Toolkit can directly read PDB files!\n",
    "topology = Topology.from_pdb(pdb_path, unique_molecules=molecules)\n",
    "\n",
    "# Construct the Interchange with the OpenFF \"Sage\" force field\n",
    "interchange = ForceField(\"openff-2.0.0.offxml\").create_interchange(topology)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tada! A beautiful solvent system:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "interchange.visualize()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "## Run a simulation\n",
    "\n",
    "We need OpenMM `System` and `Topology` objects to run our simulation, as well as positions, so lets prepare them first. We could just reuse the positions from the PDBFile and not have to worry about the units here, but in case you got your positions from somewhere else here's how to do it in the general case:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's choose parameters for the simulation and use them to prepare an Integrator:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Length of the simulation.\n",
    "num_steps = 1000  # number of integration steps to run\n",
    "\n",
    "# Logging options.\n",
    "trj_freq = 10  # number of steps per written trajectory frame\n",
    "data_freq = 10  # number of steps per written simulation statistics\n",
    "\n",
    "# Integration options\n",
    "time_step = 2 * openmm.unit.femtoseconds  # simulation timestep\n",
    "temperature = 300 * openmm.unit.kelvin  # simulation temperature\n",
    "friction = 1 / openmm.unit.picosecond  # friction constant\n",
    "\n",
    "integrator = openmm.LangevinMiddleIntegrator(temperature, friction, time_step)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Put the parts together, specify initial conditions, and configure how the simulation is recorded:"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Configure how the simulation is recorded:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "simulation = interchange.to_openmm_simulation(integrator=integrator)\n",
    "\n",
    "simulation.context.setVelocitiesToTemperature(temperature)\n",
    "\n",
    "pdb_reporter = openmm.app.PDBReporter(\"trajectory.pdb\", trj_freq)\n",
    "state_data_reporter = openmm.app.StateDataReporter(\n",
    "    \"data.csv\",\n",
    "    data_freq,\n",
    "    step=True,\n",
    "    potentialEnergy=True,\n",
    "    temperature=True,\n",
    "    density=True,\n",
    ")\n",
    "simulation.reporters.append(pdb_reporter)\n",
    "simulation.reporters.append(state_data_reporter)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We're using a PDB reporter for simplicity but you should use something more space-efficient in production. Finally, run it!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Starting simulation\")\n",
    "start = time.process_time()\n",
    "\n",
    "# Run the simulation\n",
    "simulation.step(num_steps)\n",
    "\n",
    "end = time.process_time()\n",
    "print(f\"Elapsed time {end - start} seconds\")\n",
    "print(\"Done!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can take visualize the trajectory with NGLView:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nglview.show_mdtraj(md.load(\"trajectory.pdb\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And read the produced data with Pandas:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "read_csv(\"data.csv\")"
   ]
  }
 ],
 "metadata": {
  "category": "parametrization_evaluation",
  "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}