{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Load a system from OpenMM into Interchange"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1",
   "metadata": {},
   "source": [
    "Interchange is intended as a bidirectional system conversion tool, but it's functionality for loading systems from software other than the OpenFF stack is experimental and in active development. This example describes creating an `Interchange` from an OpenMM `System` and `Topology`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pathlib\n",
    "import urllib.request\n",
    "\n",
    "import openmm\n",
    "import openmm.app\n",
    "import openmm.unit\n",
    "from openff.units.openmm import ensure_quantity\n",
    "\n",
    "from openff.interchange import Interchange\n",
    "from openff.interchange.drivers import get_openmm_energies\n",
    "from openff.interchange.drivers.openmm import _get_openmm_energies, _process"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3",
   "metadata": {},
   "source": [
    "First we'll create an OpenMM system and topology to experiment on:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def openmm_pathway(\n",
    "    pdb_file: openmm.app.PDBFile,\n",
    "    force_field: openmm.app.ForceField,\n",
    ") -> tuple[\n",
    "    openmm.System,\n",
    "    openmm.unit.Quantity,\n",
    "    openmm.unit.Quantity,\n",
    "]:\n",
    "    \"\"\"\n",
    "    Return an openmm.System, coordinates, and box vectors that result from applying this force field\n",
    "    to this PDB file.\n",
    "    \"\"\"\n",
    "    from openmm import unit\n",
    "\n",
    "    system = force_field.createSystem(\n",
    "        pdb_file.topology,\n",
    "        nonbondedCutoff=unit.Quantity(0.9, unit.nanometer),\n",
    "    )\n",
    "\n",
    "    return system, pdb_file.getPositions(), pdb_file.topology.getPeriodicBoxVectors()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not pathlib.Path(\"protein.pdb\").is_file():\n",
    "    urllib.request.urlretrieve(\n",
    "        \"https://raw.githubusercontent.com/bayer-science-for-a-better-life/abfe-benchmark/a3b15632d2f419857e2ba73a6922de6f09a66caa/structures/brd4/protein.pdb\",\n",
    "        \"brd4.pdb\",\n",
    "    )\n",
    "\n",
    "protein = openmm.app.PDBFile(\"brd4.pdb\")\n",
    "amber_force_fields = openmm.app.ForceField(\"amber14-all.xml\", \"amber14/tip3pfb.xml\")\n",
    "\n",
    "openmm_system = amber_force_fields.createSystem(\n",
    "    topology=protein.topology,\n",
    "    nonbondedMethod=openmm.app.PME,\n",
    "    nonbondedCutoff=9.0 * openmm.unit.angstrom,\n",
    "    switchDistance=8.0 * openmm.unit.angstrom,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6",
   "metadata": {},
   "source": [
    "Then, load it into Interchange. Note that the topology mentioned here is an OpenMM topology, **not** an OpenFF one."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7",
   "metadata": {},
   "outputs": [],
   "source": [
    "converted_interchange = Interchange.from_openmm(\n",
    "    topology=protein.topology,\n",
    "    system=openmm_system,\n",
    "    positions=ensure_quantity(protein.getPositions(), \"openff\"),\n",
    "    box_vectors=ensure_quantity(protein.topology.getPeriodicBoxVectors(), \"openff\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "As a partial validation of `Interchange.from_openmm`, we can compute the single-point energy of this configuration as generated by each code path. We use a high-level `get_openmm_energies` function that takes in only an `Interchange` object to get the energy of `converted_interchange` and a private `_get_openmm_energies` function to process the OpenMM objects generated by the OpenMM code path."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "converted_energies = get_openmm_energies(converted_interchange, combine_nonbonded_forces=True)\n",
    "print(converted_energies)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10",
   "metadata": {},
   "outputs": [],
   "source": [
    "(\n",
    "    reference_system,\n",
    "    reference_positions,\n",
    "    reference_box_vectors,\n",
    ") = openmm_pathway(protein, amber_force_fields)\n",
    "\n",
    "reference_energy = _process(\n",
    "    _get_openmm_energies(\n",
    "        system=reference_system,\n",
    "        positions=reference_positions,\n",
    "        box_vectors=reference_box_vectors,\n",
    "        round_positions=None,\n",
    "        platform=\"Reference\",\n",
    "    ),\n",
    "    system=reference_system,\n",
    "    combine_nonbonded_forces=True,\n",
    "    detailed=False,\n",
    ")\n",
    "\n",
    "print(reference_energy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11",
   "metadata": {},
   "source": [
    "The non-bonded energies differ slightly. It's not currently clear if this is due to difference in non-bonded settings or a subtle difference in the conversion of applied parameters. Energies from valence terms match to a high level of precision."
   ]
  }
 ],
 "metadata": {
  "category": "force_field_interop",
  "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
}