section_3/machine_learning_potentials.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "HL9clVbBmMu-"
      },
      "source": [
        "# Machine Learning Potentials\n",
        "\n",
        "[![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/sef43/openmm_workshop/blob/main/section_3/machine_learning_potentials.ipynb)\n",
        "\n",
        "## Table of contents\n",
        "- [Introduction](#intro)\n",
        "- [OpenMM ML software](#openmmmlsoftware)\n",
        "- [Basics of OpenMM compatible MLPs](#basicsofmlps)\n",
        "- [Installing](#installing)\n",
        "- [Exporting a PyTorch model for use in OpenMM](#pytorchmodelinopenmm)\n",
        "- [Simulation of alanine dipeptide with ANI-2x using OpenMM-Torch](#ani)\n",
        "    - [Prepate test system](#prepare)\n",
        "    - [Create a NNP](#nnp)\n",
        "    - [Adding the NNP to an OpenMM simulation](#addnnp)\n",
        "- [Mixed MM/ML system](#mixedsystem)\n",
        "    - [Creating the system](#createmixed)\n",
        "    - [create the MLP](#createmlp)\n",
        "- [Using OpenMM-ML package](#openmmml)\n",
        "- [Using NNPOps](#nnpops)\n",
        "- [Implementing other models - MACE](#mace)\n",
        "- [Extra exercises](#extra)\n",
        "- [References](#references)\n",
        "- [Solutions](#solutions)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "## Introduction\n",
        "<a id=\"intro\"></a>\n",
        "\n",
        "Machine Learning Potentials (MLPs) are a relatively new method where the potential energy surface of an atomic system is described by some sort of machine learning model. The model could be a feed forward Neural Network (ANI), Gaussian process regression (GAP), Graph neural network (MACE), Equivariant Transformer (TorchMD-NET), or something else. The MLP is trained on a first principles method such as DFT. It can then be used in an atomistic MD simulation in the place of a classical forcefield bringing the accuracy of first principles methods without the computational expense."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "## OpenMM ML software\n",
        "<a id=\"openmmmlsoftware\"></a>\n",
        "\n",
        "OpenMM is an MD engine therefore we will be covering how to use MLPs in a simulation and not cover how to train them.\n",
        "\n",
        "OpenMM has serval packages supporting the use of MLPs.\n",
        "- [openmm-torch](https://github.com/openmm/openmm-torch). This is the OpenMM PyTorch plugin what allows [PyTorch](https://pytorch.org/) static computation graphs to be used for defining an OpenMM TorchForce object, an [OpenMM Force class](http://docs.openmm.org/latest/api-python/library.html#forces) that computes a contribution to the potential energy.\n",
        "- [openmm-ml](https://github.com/openmm/openmm-ml). This is a high level API for using machine learning models in OpenMM simulations. With just a few lines of code, you can set up a simulation that uses a standard, pretrained model to represent some or all of the interactions in a system.\n",
        "- [NNPops](https://github.com/openmm/NNPOps). This is a library of optimized operations that appear in popular Neural Network Potentials that can be used to speed up your MLP implementation.\n",
        "\n",
        "\n",
        "You would use `openmm-torch` to interface your MLP with OpenMM. You could then try and optimize it using operations from `NNPOps`. When you want to deploy it you can use `openmm-ml` to create an easy to use wrapper around it."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Basics of (OpenMM compatible) MLPs\n",
        "<a id=\"basicsofmlps\"></a>\n",
        "\n",
        "A MLP (or NNP, we will use these terms interchangeably) reads in a set of atomic coordinates and outputs the potential energy. The forces on each particle can then be computed by backpropagation and taking the gradient of the energy with respect to the coordinates ($F=-\\nabla V$).\n",
        "\n",
        "To use a MLP in OpenMM you need to be able to write it as a PyTorch model than can be exported to [TorchScript]([https://pytorch.org/docs/stable/jit.html#torchscript]). The model takes a (nparticle, 3) shape tensor of particle positions and produces the energy. openmm-torch will calculate the forces. Optionally the model can return the forces and openmm-torch will use them directly."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Installing packages\n",
        "<a id=\"installing\"></a>\n",
        "\n",
        "The packages can be installed from conda-forge. Note that there is no windows package for `openmm-torch` and there are only linux packages for `NNPOps`.\n",
        "\n",
        "We will install `openmm-torch` and at the same time install `torch-ani` which we will need later."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "9aRhmqIz7-2q",
        "outputId": "93ec5ae8-7cf2-4519-828d-fec21d316c41"
      },
      "outputs": [],
      "source": [
        "# Execute this cell to install mamba in the Colab environment\n",
        "\n",
        "if 'google.colab' in str(get_ipython()):\n",
        "  print('Running on colab')\n",
        "  !pip install -q condacolab\n",
        "  import condacolab\n",
        "  condacolab.install_mambaforge()\n",
        "else:\n",
        "  print('Not running on colab.')\n",
        "  print('Make sure you create and activate a new conda environment!')\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W0RghZ91mMvC",
        "outputId": "623e0a70-fba2-459a-99f1-a69c7c0cfebf"
      },
      "outputs": [],
      "source": [
        "if 'google.colab' in str(get_ipython()):\n",
        "  #https://github.com/openmm/openmm-torch/issues/88\n",
        "  %env CONDA_OVERRIDE_CUDA=12.0\n",
        "  # you might also need to set this is you are on linux without CUDA installed!\n",
        "!mamba install -y -c conda-forge openmm-torch torchani=2.2.2"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Download the files we will need."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!wget https://raw.githubusercontent.com/sef43/openmm_workshop/main/section_3/alanine-dipeptide.pdb\n",
        "!wget https://raw.githubusercontent.com/sef43/openmm_workshop/main/section_3/section_3_utils.py"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "bYPvb1EomMvD"
      },
      "source": [
        "## Exporting a PyTorch model for use in OpenMM\n",
        "<a id=\"pytorchmodelinopenmm\"></a>\n",
        "\n",
        "We can check that our installation is working by defining a very simple potential --- a harmonic force attracting every particle to the origin.\n",
        "\n",
        "The first step is to create a PyTorch model defining the calculation. It should take the particle positions in nanometers (in the form of `torch.Tensor` of shape `(nparticles, 3)` as input, and return the potential energy in kJ/mol.)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LVXLNA5RmMvE"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "class ForceModule(torch.nn.Module):\n",
        "    \"\"\"A central harmonic potential as a static compute graph\"\"\"\n",
        "    def forward(self, positions: torch.Tensor):\n",
        "        \"\"\"The forward method returns the energy computed from positions.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        positions : torch.Tensor with shape (nparticles,3)\n",
        "           positions[i,k] is the position (in nanometers) of spatial dimension k of particle i\n",
        "\n",
        "        Returns\n",
        "        -------\n",
        "        potential : torch.Tensor\n",
        "           The potential energy (in kJ/mol)\n",
        "        \"\"\"\n",
        "        return torch.sum(positions**2)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "XnUjIBBAmMvE"
      },
      "source": [
        "The ForceModule inherits from [`torch.nn.module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) which is a pytorch base class for neural network modules. We put the code in the `forward` method. The forward method defines the computation performed at every call of the model."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-LYilcYPmMvE"
      },
      "source": [
        "\n",
        "The model must then be converted to a TorchScript module and saved to a file. Converting to TorchScript can usually be done with a single call to torch.jit.script. See the [PyTorch documentation](https://pytorch.org/tutorials/beginner/Intro_to_TorchScript_tutorial.html) for details.  "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7kSTrSMZmMvF"
      },
      "outputs": [],
      "source": [
        "# convert to TorchScript\n",
        "module = torch.jit.script(ForceModule())"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "QvYVFsYYmMvF"
      },
      "source": [
        "We then need to save the module. The saving process serializes the module. This means we can load it into the C++ API as required by openmm-torch `TorchForce`. See the [PyTorch docs](https://pytorch.org/docs/stable/generated/torch.jit.save.html) for more information."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lK91t53WmMvF"
      },
      "outputs": [],
      "source": [
        "# saved the serialized compute graph to a file\n",
        "module.save('model.pt')"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "0bNpF1rpmMvG"
      },
      "source": [
        "To use the exported model in a simulation we need to create a TorchForce object and add it to the OpenMM System."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "6aTwPWyOmMvG",
        "outputId": "d695198a-2de7-4153-b1aa-6bcf2ff0c825"
      },
      "outputs": [],
      "source": [
        "# Create the TorchForce from the serialized compute graph\n",
        "from openmmtorch import TorchForce\n",
        "torch_force = TorchForce('model.pt')\n",
        "\n",
        "# Create an empty OpenMM system\n",
        "import openmm as mm\n",
        "system = mm.System()\n",
        "\n",
        "# add the TorchForce to the system\n",
        "system.addForce(torch_force)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-Ym2SuXPmMvG"
      },
      "source": [
        "Now that we know how to use a PyTorch model in an OpenMM simulation we will move into an example of creating a Neural Network Potential that uses ANI-2x."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n",
        "## Simulation of alanine dipeptide with ANI-2x using OpenMM-Torch\n",
        "<a id=\"ani\"></a>\n",
        "\n",
        "ANI-2x is a general Neural Network Potential that works with molecules containing (H, C, N, O, F, Cl, S) atoms. For more information please read the publication.\n",
        "The model is available from the TorchANI package which we installed earlier."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ALbPQnXdmMvG"
      },
      "source": [
        "### Prepare a test system\n",
        "<a id=\"prepare\"></a>\n",
        "\n",
        "For simplicity we will use an alanine-dipeptide test system. We prepare it as we have done before but then we remove all the standard MM forces."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wOaWuGs9mMvG"
      },
      "outputs": [],
      "source": [
        "import openmm as mm\n",
        "import openmm.app as app\n",
        "import openmm.unit as unit\n",
        "\n",
        "# create an alanine-dipeptide test system\n",
        "pdb = app.PDBFile('alanine-dipeptide.pdb')\n",
        "forcefield = app.ForceField('amber14-all.xml')\n",
        "system = forcefield.createSystem(pdb.topology, constraints=None)\n",
        "\n",
        "# Remove MM forces\n",
        "while system.getNumForces() > 0:\n",
        "  system.removeForce(0)\n",
        "\n",
        "# The system should not contain any additional force and constrains\n",
        "assert system.getNumConstraints() == 0\n",
        "assert system.getNumForces() == 0\n",
        "\n",
        "\n",
        "# Get the list of atomic numbers. We will need this when creating the model instance\n",
        "atomic_numbers = [atom.element.atomic_number for atom in pdb.topology.atoms()]"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "agpWvioxmMvG"
      },
      "source": [
        "### Define the NNP\n",
        "<a id=\"nnp\"></a>\n",
        "\n",
        "We can create a NNP class that uses ANI-2x."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Xhk7OGiEmMvG",
        "outputId": "94760d61-19db-485a-f1de-2e5d5f24d3bc"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torchani.models import ANI2x\n",
        "\n",
        "class NNP(torch.nn.Module):\n",
        "\n",
        "  def __init__(self, atomic_numbers: torch.Tensor):\n",
        "\n",
        "    super().__init__()\n",
        "\n",
        "    # Store the atomic numbers\n",
        "    self.atomic_numbers = atomic_numbers.unsqueeze(0)\n",
        "\n",
        "    # Create an ANI-2x model\n",
        "    self.model = ANI2x(periodic_table_index=True)\n",
        "\n",
        "    # make sure it is on the same device at the atomic_numbers tensor\n",
        "    self.model.to(self.atomic_numbers.device)\n",
        "\n",
        "  def forward(self, positions: torch.Tensor):\n",
        "\n",
        "    # Prepare the positions\n",
        "    positions = positions.unsqueeze(0).float() * 10 # nm --> Å\n",
        "\n",
        "    # Run ANI-2x\n",
        "    result = self.model((self.atomic_numbers, positions))\n",
        "\n",
        "    # Get the potential energy\n",
        "    energy = result.energies[0] * 2625.5 # Hartree --> kJ/mol\n",
        "\n",
        "    return energy\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "18qSwMKJ57Qs"
      },
      "source": [
        "### Create NNP\n",
        "<a id=\"creatennp\"></a>\n",
        "\n",
        "We can now create an instance of the NNP making sure to use the gpu ('cuda' device) if available."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D70al7eg5w6Y",
        "outputId": "489bf0f5-9c40-4b36-bedf-946923de9626"
      },
      "outputs": [],
      "source": [
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "nnp = NNP(torch.tensor(atomic_numbers,device=device))\n",
        "\n",
        "#https://github.com/aiqm/torchani/issues/628\n",
        "torch._C._jit_set_nvfuser_enabled(False)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "mMQ2u0JGmMvH"
      },
      "source": [
        "We can now compute the potential energy of the system using the PyTorch interface. We can also compute the forces using autograd."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "Q049ZnK5mMvH",
        "outputId": "7c44b7d9-7b89-4351-aa2b-7f634034a442"
      },
      "outputs": [],
      "source": [
        "# Need to make a torch.tensor of the positions. Require grad so we can compute the forces.\n",
        "positions = torch.tensor(pdb.positions.value_in_unit(unit.nanometers), device=device, requires_grad=True)\n",
        "\n",
        "# put the positions into the NNP and it returns the energy\n",
        "energy = nnp(positions)\n",
        "\n",
        "print(energy)\n",
        "\n",
        "# we can compute the forces using autograd\n",
        "energy.backward()\n",
        "force = -positions.grad\n",
        "\n",
        "print(force)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "bByhu-_YmMvH"
      },
      "source": [
        "### Add the NNP to the system\n",
        "<a id=\"addnnp\"></a>\n",
        "\n",
        "We now export the model and load it with `TorchForce`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "64NJgUdDmMvH",
        "outputId": "6a43d17d-5a3c-4770-eb34-0cefe7d1624c"
      },
      "outputs": [],
      "source": [
        "from openmmtorch import TorchForce\n",
        "import sys\n",
        "\n",
        "# Save the NNP to a file and load it with OpenMM-Torch\n",
        "torch.jit.script(nnp).save('model.pt')\n",
        "torchforce = TorchForce('model.pt')\n",
        "\n",
        "# Add the NNP to the system\n",
        "system.addForce(torchforce)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "PR7qDA5VmMvH"
      },
      "source": [
        "### Create a simulation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5E4dOB1CmMvH"
      },
      "outputs": [],
      "source": [
        "# Create an integrator with a time step of 0.5 fs\n",
        "temperature = 298.15 * unit.kelvin\n",
        "frictionCoeff = 1 / unit.picosecond\n",
        "timeStep = 0.5 * unit.femtosecond\n",
        "integrator = mm.LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)\n",
        "\n",
        "# Create a simulation and set the initial positions and velocities\n",
        "simulation = app.Simulation(pdb.topology, system, integrator)\n",
        "simulation.context.setPositions(pdb.positions)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Z9cJiI0xmMvH"
      },
      "source": [
        "Now we can compute the energy and forces again but using the OpenMM interface. We compare them to the energy and forces computed from the PyTorch interface. This is a good check to do because sometimes bugs can arise during the export/serialization and load step."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CkFPCAiUmMvH",
        "outputId": "364ed990-da27-4a65-ae01-0d012c235bb7"
      },
      "outputs": [],
      "source": [
        "state = simulation.context.getState(getEnergy=True, getForces=True)\n",
        "openmm_energy = state.getPotentialEnergy().value_in_unit(unit.kilojoule_per_mole)\n",
        "openmm_force  =  state.getForces(asNumpy=True).value_in_unit(unit.kilojoule_per_mole/unit.nanometer)\n",
        "\n",
        "print(openmm_energy)\n",
        "print(openmm_force)\n",
        "import numpy as np\n",
        "\n",
        "assert(np.isclose(openmm_energy, energy.cpu().detach().numpy()))\n",
        "assert(np.allclose(openmm_force, force.cpu().detach().numpy(),rtol=1e-3))"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "1p34mapBmMvH"
      },
      "source": [
        "### Run the simulation"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ZKBxSO8mmMvH",
        "outputId": "eab2f8a5-68d1-4d70-85cf-ca73440a8fb6"
      },
      "outputs": [],
      "source": [
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ye5bo0PVavZI"
      },
      "source": [
        "## Mixed system\n",
        "<a id=\"mixedsystem\"></a>\n",
        "\n",
        "In general ML forecfields are still too computationally slow to be used to model entire solvated biomolecules. However, a use-case to exploit their accuracy without the prohibitive cost is to model a small part of the system with the MLP and the rest of the system with a traditional MM forcefield [[1, 2]](#references).\n",
        "For example a ligand's intramolecular energy could be modelled with the MLP and the rest of the system, including intermolecular interactions between ligand and the protein/solvent, would be modelled with a MM forecfield. This approach is similar to hybrid QM/MM methods.\n",
        "\n",
        "Our example system will be the alanine-dipeptide in a water box.\n",
        "We will model the alanine-dipeptide intramolecular forces with the ANI-2x MLP and the water molecules will use a MM forcefield, addiontially the intermolecular interactions between the alanine-dipeptide and the water will use the MM forcefield.\n",
        "\n",
        "The total potential energy of the mixed system can be written as\n",
        "$V_{MM/ML}(r) = V_{MM}(r_{MM}) + V_{MM-ML}(r) + V_{ML}(r_{ML})$ ,\n",
        "where $r$ are the coordinates of all atoms, $r_{MM}$ are the coordinates of just the atoms in the MM region, and $r_{ML}$ are the coordinates of just the atoms in the ML region. The three terms are:\n",
        "\n",
        "  - $V_{MM}(r_{MM})$ - The potential energy of the MM region (water molecules) using the MM forcefield.\n",
        "  - $V_{MM-ML}(r)$ - The coupling term between the MM and ML regions. We will define this to compute the non-bonded intermolecular interactions between the ML region and MM region atoms using the MM forcefield.\n",
        "  - $V_{ML}(r_{ML})$ - The intramolecular potential energy of the ML region (alanine-dipeptide) using the ML forcefield."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Creating the system\n",
        "<a id=\"createmixed\"></a>\n",
        "\n",
        "We will first create a MM system as normal. Then we will define a function that modifies it to use the hybrid potential described above."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rx66FzALH_Ql"
      },
      "outputs": [],
      "source": [
        "import openmm as mm\n",
        "import openmm.app as app\n",
        "import openmm.unit as unit\n",
        "\n",
        "# create an alanine-dipeptide test system\n",
        "pdb = app.PDBFile('alanine-dipeptide.pdb')\n",
        "forcefield = app.ForceField('amber14-all.xml', 'amber14/tip3p.xml')\n",
        "modeller = app.Modeller(pdb.topology, pdb.positions)\n",
        "modeller.addSolvent(forcefield, padding=1.0*unit.nanometers)\n",
        "system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.PME, constraints=None)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Y98CrfILMfIM"
      },
      "source": [
        "Now we create a function that will remove the MM interactions within the ML region."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "3HXL44jVJnND"
      },
      "outputs": [],
      "source": [
        "from section_3_utils import removeBonds\n",
        "\n",
        "def removeMMInteraction(topology, system, ml_atoms):\n",
        "\n",
        "  # remove the bonded interactions witin the ML subset\n",
        "  newSystem = removeBonds(system, ml_atoms)\n",
        "\n",
        "  # Add nonbonded exceptions and exclusions.\n",
        "  # This removes the nonbonded interactions between the ML atoms\n",
        "  atomList = list(ml_atoms)\n",
        "  for force in newSystem.getForces():\n",
        "      if isinstance(force, mm.NonbondedForce):\n",
        "          for i in range(len(atomList)):\n",
        "              for j in range(i):\n",
        "                  force.addException(i, j, 0, 1, 0, True)\n",
        "      elif isinstance(force, mm.CustomNonbondedForce):\n",
        "          existing = set(tuple(force.getExclusionParticles(i)) for i in range(force.getNumExclusions()))\n",
        "          for i in range(len(atomList)):\n",
        "              for j in range(i):\n",
        "                  if (i, j) not in existing and (j, i) not in existing:\n",
        "                      force.addExclusion(i, j, True)\n",
        "\n",
        "  return newSystem\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "4Uc6bgKuMxoF"
      },
      "source": [
        "### Create the MLP for a mixed system\n",
        "<a id=\"createmlp\"></a>\n",
        "\n",
        "We will create an ANI-2x MLP as before but add in an extra argument that lists the atoms in the ML region."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "OIdEnzRFNIAY"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from torchani.models import ANI2x\n",
        "\n",
        "class hybridNNP(torch.nn.Module):\n",
        "\n",
        "  def __init__(self, atomic_numbers: torch.Tensor, ml_atoms: torch.Tensor):\n",
        "\n",
        "    super().__init__()\n",
        "\n",
        "    # The atomic_numbers tensor contains the atomic number of the ML region atoms.\n",
        "    # the ml_atoms tensor contains the index of each ML atom with respect to the full system.\n",
        "    assert(atomic_numbers.shape == ml_atoms.shape)\n",
        "\n",
        "    # Store the indices of the ml atoms\n",
        "    self.indices = ml_atoms\n",
        "\n",
        "    # Store the atomic numbers\n",
        "    self.atomic_numbers = atomic_numbers.unsqueeze(0)\n",
        "\n",
        "    # Create an ANI-2x model\n",
        "    self.model = ANI2x(periodic_table_index=True)\n",
        "\n",
        "    # make sure it is on the same device at the atomic_numbers tensor\n",
        "    self.model.to(self.atomic_numbers.device)\n",
        "\n",
        "  def forward(self, positions: torch.Tensor):\n",
        "\n",
        "    # extract the positions of the ML atoms\n",
        "    positions = positions[self.indices]\n",
        "\n",
        "    # Prepare the positions\n",
        "    positions = positions.unsqueeze(0).float() * 10 # nm --> Å\n",
        "\n",
        "    # Run ANI-2x\n",
        "    result = self.model((self.atomic_numbers, positions))\n",
        "\n",
        "    # Get the potential energy\n",
        "    energy = result.energies[0] * 2625.5 # Hartree --> kJ/mol\n",
        "\n",
        "    return energy\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "CkdqLIvCPRAL"
      },
      "source": [
        "Now we can create an instance of the MLP and add it to the system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "H8SgRpvDPbpt",
        "outputId": "549fdb28-6e6b-4102-a455-044c17c8e109"
      },
      "outputs": [],
      "source": [
        "# get a list of the ML atoms. The alanine-dipeptide is chain 0.\n",
        "chains = list(modeller.topology.chains())\n",
        "ml_atoms = [atom.index for atom in chains[0].atoms()]\n",
        "print(ml_atoms)\n",
        "\n",
        "# get the atomic numbers\n",
        "atomic_numbers = [atom.element.atomic_number for atom in chains[0].atoms()]\n",
        "\n",
        "# convert to torch tensors\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "ml_atoms = torch.tensor(ml_atoms, device=device, dtype=torch.int64)\n",
        "atomic_numbers = torch.tensor(atomic_numbers, device=device)\n",
        "\n",
        "hybridnnp = hybridNNP(atomic_numbers, ml_atoms)\n",
        "\n",
        "#https://github.com/aiqm/torchani/issues/628\n",
        "torch._C._jit_set_nvfuser_enabled(False)\n",
        "\n",
        "# Save the NNP to a file and load it with OpenMM-Torch\n",
        "torch.jit.script(hybridnnp).save('mixed_model.pt')\n",
        "\n",
        "from openmmtorch import TorchForce\n",
        "torchforce = TorchForce('mixed_model.pt')\n",
        "\n",
        "# make the mixed system\n",
        "mixed_system = removeMMInteraction(modeller.topology, system, ml_atoms.tolist())\n",
        "\n",
        "# add the TorchForce\n",
        "mixed_system.addForce(torchforce)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "taMDOGCpUaqd"
      },
      "source": [
        "### Simulate"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "IIMkyPtXUh83",
        "outputId": "62638193-5a17-49fe-aee5-fcbe728c5dba"
      },
      "outputs": [],
      "source": [
        "# Create an integrator with a time step of 0.5 fs\n",
        "temperature = 298.15 * unit.kelvin\n",
        "frictionCoeff = 1 / unit.picosecond\n",
        "timeStep = 0.5 * unit.femtosecond\n",
        "integrator = mm.LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)\n",
        "\n",
        "# Create a simulation and set the initial positions and velocities\n",
        "simulation = app.Simulation(modeller.topology, mixed_system, integrator)\n",
        "simulation.context.setPositions(modeller.positions)\n",
        "\n",
        "simulation.minimizeEnergy(maxIterations=100)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)\n",
        "\n",
        "simulation.reporters.append(app.PDBReporter('mixed_traj.pdb', 100))\n",
        "\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "wkd_ANeLZTZ_"
      },
      "source": [
        "## Using the Openmm-ML package\n",
        "<a id=\"openmmml\"></a>\n",
        "\n",
        "We have covered how to use openmm-torch to add a MLP to a system. As you have seen it involves quite a lot of code. The Openmm-ML package was created to be a high level interface for people to use pre-trained ML models in their OpenMM simulations. We will now do the same simulations using openmm-ml\n",
        "\n",
        "### Install software\n",
        "The openmm-ml package can be installed from [conda-forge](https://anaconda.org/conda-forge/openmm-ml)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "bVooMbtCafwB",
        "outputId": "f0c55a44-4281-48d6-8aca-3355d42d98c6"
      },
      "outputs": [],
      "source": [
        "!mamba install -y -c conda-forge openmm-ml"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "6z5bqtC2dGxC"
      },
      "source": [
        "### Create a pure ML system\n",
        "\n",
        "We will load in the alanine-dipeptide molecule and simulate it in vacuum with ANI-2x using the OpenMM-ML `MLPPotential.createSystem` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "W1s4hdDZc155",
        "outputId": "d27166c0-a15f-409d-93f8-dc9e01be244b"
      },
      "outputs": [],
      "source": [
        "import openmm as mm\n",
        "import openmm.app as app\n",
        "import openmm.unit as unit\n",
        "\n",
        "from openmmml import MLPotential\n",
        "\n",
        "pdb = app.PDBFile('alanine-dipeptide.pdb')\n",
        "\n",
        "# create the MLP using ANI-2x\n",
        "potential = MLPotential('ani2x')\n",
        "\n",
        "# create a system that uses the MLP\n",
        "system = potential.createSystem(pdb.topology)\n",
        "\n",
        "\n",
        "# Create an integrator with a time step of 0.5 fs\n",
        "temperature = 298.15 * unit.kelvin\n",
        "frictionCoeff = 1 / unit.picosecond\n",
        "timeStep = 0.5 * unit.femtosecond\n",
        "integrator = mm.LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)\n",
        "\n",
        "# Create a simulation and set the initial positions and velocities\n",
        "simulation = app.Simulation(pdb.topology, system, integrator)\n",
        "simulation.context.setPositions(pdb.positions)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)\n",
        "\n",
        "simulation.step(1000)\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "You can safely ignore the error message that says `failed to equip 'nnpops' with error: No module named 'NNPOps'`.\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "0VIiA4tSeTBJ"
      },
      "source": [
        "### Create a mixed system\n",
        "We can just as easily create a mixed system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "HdtTiERveYv8",
        "outputId": "75d05af5-f8c8-4bf8-b1cf-7a31d7453bfc"
      },
      "outputs": [],
      "source": [
        "import openmm as mm\n",
        "import openmm.app as app\n",
        "import openmm.unit as unit\n",
        "\n",
        "from openmmml import MLPotential\n",
        "\n",
        "pdb = app.PDBFile('alanine-dipeptide.pdb')\n",
        "\n",
        "forcefield = app.ForceField('amber14-all.xml', 'amber14/tip3p.xml')\n",
        "modeller = app.Modeller(pdb.topology, pdb.positions)\n",
        "modeller.addSolvent(forcefield, padding=1.0*unit.nanometers)\n",
        "\n",
        "\n",
        "# create the MM system\n",
        "mm_system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.PME, constraints=None)\n",
        "\n",
        "# create the MLP using ANI-2x\n",
        "potential = MLPotential('ani2x')\n",
        "\n",
        "\n",
        "# create the mixed system\n",
        "chains = list(modeller.topology.chains())\n",
        "ml_atoms = [atom.index for atom in chains[0].atoms()]\n",
        "potential = MLPotential('ani2x')\n",
        "mixed_system = potential.createMixedSystem(modeller.topology, mm_system, ml_atoms)\n",
        "\n",
        "\n",
        "# Create an integrator with a time step of 0.5 fs\n",
        "temperature = 298.15 * unit.kelvin\n",
        "frictionCoeff = 1 / unit.picosecond\n",
        "timeStep = 0.5 * unit.femtosecond\n",
        "integrator = mm.LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)\n",
        "\n",
        "# Create a simulation and set the initial positions and velocities\n",
        "simulation = app.Simulation(modeller.topology, mixed_system, integrator)\n",
        "simulation.context.setPositions(modeller.positions)\n",
        "simulation.minimizeEnergy(maxIterations=100)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)\n",
        "\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "SqlNOUd-gEfZ"
      },
      "source": [
        "## Using NNPOps\n",
        "<a id=\"nnpops\"></a>\n",
        "\n",
        "The NNPOps package provides highly optimized, open source implementations of bottleneck operations that appear in popular potentials. It can be used to speed up ANI simulations. We can install it from conda-forge (only on Linux) and use it through the Openmm-ml interface."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "1yT_Htmqgldx",
        "outputId": "3b1a9208-6f02-4280-a398-d6154ae5ab3e"
      },
      "outputs": [],
      "source": [
        "!mamba install -y -c conda-forge nnpops"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "foXNxdQ4gtRZ"
      },
      "source": [
        "If you are using a GPU it should offer you some speed up. We will use the script from before and specifiy to openmm-ml that is should use NNPOps."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D3DIUf5Ugs_N",
        "outputId": "d8ac6278-df38-49c6-b4fd-221294913d0a"
      },
      "outputs": [],
      "source": [
        "import openmm as mm\n",
        "import openmm.app as app\n",
        "import openmm.unit as unit\n",
        "\n",
        "from openmmml import MLPotential\n",
        "\n",
        "pdb = app.PDBFile('alanine-dipeptide.pdb')\n",
        "\n",
        "# create the MLP using ANI-2x\n",
        "# Use NNPOps\n",
        "potential = MLPotential('ani2x', implementation='nnpops')\n",
        "\n",
        "# create a system that uses the MLP\n",
        "system = potential.createSystem(pdb.topology)\n",
        "\n",
        "\n",
        "# Create an integrator with a time step of 0.5 fs\n",
        "temperature = 298.15 * unit.kelvin\n",
        "frictionCoeff = 1 / unit.picosecond\n",
        "timeStep = 0.5 * unit.femtosecond\n",
        "integrator = mm.LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)\n",
        "\n",
        "# Create a simulation and set the initial positions and velocities\n",
        "simulation = app.Simulation(pdb.topology, system, integrator)\n",
        "simulation.context.setPositions(pdb.positions)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)\n",
        "\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "zBe8H9m5hP6A"
      },
      "source": [
        "The `implementation='nnpops'` argument tells 'MLPotential that it should swap out some of the torchANI pytorch functions for the optimized NNPOps versions. If you want to see how this is done in the code take a look at the [github repo](https://github.com/openmm/NNPOps)."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "aif9Zl-Zhxup"
      },
      "source": [
        "## Implementing other Models - MACE\n",
        "<a id=\"mace\"></a>\n",
        "\n",
        "\n",
        "Any model that can be written as a pytorch model can be used with OpenMM-torch. We will show how to implement a MACE model.\n",
        "\n",
        "Please also take a look at the [MACE documentation](https://mace-docs.readthedocs.io/en/latest/guide/openmm.html) about their `mace-md` OpenMM interface.\n",
        "\n",
        "### Install software\n",
        "MACE can be installed with pip."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "txm22xKliRRl"
      },
      "outputs": [],
      "source": [
        "!pip install git+https://github.com/ACEsuit/mace"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Ong6DOIzia1r"
      },
      "source": [
        "### Get a pretrained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "anI4ZqDyjD5J",
        "outputId": "d74bdd41-f3bd-4152-fc7c-ee050aafae81"
      },
      "outputs": [],
      "source": [
        "!wget https://github.com/ACEsuit/mace/raw/docs/docs/examples/ANI_trained_MACE.zip\n",
        "!unzip ANI_trained_MACE.zip"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "EE9DpkY7jmtK"
      },
      "source": [
        "### Define the MLP\n",
        "We will create a MACE MLP class as we did before for ANI-2x. The MACE model requires some extra code to covert from atomic numbers and positions into the required format. Some of this code has been put in the section_3_utils.py module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Y2EBX4TkNdl"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from section_3_utils import simple_nl\n",
        "from e3nn.util import jit\n",
        "from mace.tools import utils, to_one_hot, atomic_numbers_to_indices\n",
        "from typing import Optional\n",
        "\n",
        "class MACEForce(torch.nn.Module):\n",
        "  def __init__(self, model_path, atomic_numbers, indices, periodic, device, dtype=torch.float64):\n",
        "      super().__init__()\n",
        "\n",
        "      if device is None: # use cuda if available\n",
        "          self.device=torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "\n",
        "      else: # unless user has specified the device\n",
        "          self.device=torch.device(device)\n",
        "\n",
        "      self.default_dtype = dtype\n",
        "      torch.set_default_dtype(self.default_dtype)\n",
        "\n",
        "      print(\"Running MACEForce on device: \", self.device, \" with dtype: \", self.default_dtype)\n",
        "\n",
        "      # conversion constants\n",
        "      self.nm_to_distance = 10.0 # nm->A\n",
        "      self.distance_to_nm = 0.1 # A->nm\n",
        "      self.energy_to_kJ = 96.49 # eV->kJ\n",
        "\n",
        "      self.model = torch.load(model_path,map_location=device)\n",
        "      self.model.to(self.default_dtype)\n",
        "      self.model.eval()\n",
        "\n",
        "\n",
        "      self.r_max = self.model.r_max\n",
        "      self.z_table = utils.AtomicNumberTable([int(z) for z in self.model.atomic_numbers])\n",
        "\n",
        "      self.model = jit.compile(self.model)\n",
        "\n",
        "      # setup input\n",
        "      N=len(atomic_numbers)\n",
        "      self.ptr = torch.tensor([0,N],dtype=torch.long, device=self.device)\n",
        "      self.batch = torch.zeros(N, dtype=torch.long, device=self.device)\n",
        "\n",
        "      # one hot encoding of atomic number\n",
        "      self.node_attrs = to_one_hot(\n",
        "              torch.tensor(atomic_numbers_to_indices(atomic_numbers, z_table=self.z_table), dtype=torch.long, device=self.device).unsqueeze(-1),\n",
        "              num_classes=len(self.z_table),\n",
        "          )\n",
        "\n",
        "      if periodic:\n",
        "          self.pbc=torch.tensor([True, True, True], device=self.device)\n",
        "      else:\n",
        "          self.pbc=torch.tensor([False, False, False], device=self.device)\n",
        "\n",
        "      if indices is None:\n",
        "          self.indices = None\n",
        "      else:\n",
        "          self.indices = torch.tensor(indices, dtype=torch.int64)\n",
        "\n",
        "\n",
        "\n",
        "  def forward(self, positions, boxvectors: Optional[torch.Tensor] = None):\n",
        "      # setup positions\n",
        "\n",
        "      positions = positions.to(device=self.device,dtype=self.default_dtype)\n",
        "      if self.indices is not None:\n",
        "          positions = positions[self.indices]\n",
        "\n",
        "      positions = positions*self.nm_to_distance\n",
        "\n",
        "      if boxvectors is not None:\n",
        "          cell = boxvectors.to(device=self.device,dtype=self.default_dtype) * self.nm_to_distance\n",
        "          pbc = True\n",
        "      else:\n",
        "          cell = torch.eye(3, device=self.device)\n",
        "          pbc = False\n",
        "\n",
        "      mapping, shifts_idx = simple_nl(positions, cell, pbc, self.r_max)\n",
        "\n",
        "      edge_index = torch.stack((mapping[0], mapping[1]))\n",
        "\n",
        "      shifts = torch.mm(shifts_idx, cell)\n",
        "\n",
        "      # create input dict\n",
        "      input_dict = { \"ptr\" : self.ptr,\n",
        "                    \"node_attrs\": self.node_attrs,\n",
        "                    \"batch\": self.batch,\n",
        "                    \"pbc\": self.pbc,\n",
        "                    \"cell\": cell,\n",
        "                    \"positions\": positions,\n",
        "                    \"edge_index\": edge_index,\n",
        "                    \"unit_shifts\": shifts_idx,\n",
        "                    \"shifts\": shifts}\n",
        "\n",
        "      # predict\n",
        "      out = self.model(input_dict,compute_force=False)\n",
        "\n",
        "      energy = out[\"interaction_energy\"]\n",
        "      if energy is None:\n",
        "          energy = torch.tensor(0.0, device=self.device)\n",
        "\n",
        "      # return energy\n",
        "      energy = energy*self.energy_to_kJ\n",
        "\n",
        "      return energy\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "IKKPko8tmrwq"
      },
      "source": [
        "### Use the MACE MLP\n",
        "\n",
        "the rest of the code is then similar to use the ANI-2x MLP via the openmm-torch interface."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 407
        },
        "id": "BbDMuwx2mrMP",
        "outputId": "e23cc163-a1be-4e64-b26e-9023d00d5621"
      },
      "outputs": [],
      "source": [
        "import openmm as mm\n",
        "import openmm.app as app\n",
        "import openmm.unit as unit\n",
        "from openmmtorch import TorchForce\n",
        "\n",
        "pdb = app.PDBFile('alanine-dipeptide.pdb')\n",
        "forcefield = app.ForceField('amber14-all.xml')\n",
        "system = forcefield.createSystem(pdb.topology, constraints=None)\n",
        "# Remove MM forces\n",
        "while system.getNumForces() > 0:\n",
        "  system.removeForce(0)\n",
        "# The system should not contain any additional force and constrains\n",
        "assert system.getNumConstraints() == 0\n",
        "assert system.getNumForces() == 0\n",
        "\n",
        "# Get the list of atomic numbers.\n",
        "atomic_numbers = [atom.element.atomic_number for atom in pdb.topology.atoms()]\n",
        "\n",
        "# Create the MACE MLP\n",
        "model_path = \"ANI_trained/ani500k_small_DFT.model\"\n",
        "pbc = False\n",
        "indices = None\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "\n",
        "mace_mlp = MACEForce(model_path, atomic_numbers, indices, pbc, device)\n",
        "\n",
        "# export it with torchscript\n",
        "torch.jit.script(mace_mlp).save('macemodel.pt')\n",
        "\n",
        "# load it in with torchforce\n",
        "torchforce = TorchForce('macemodel.pt')\n",
        "\n",
        "# add it to the system\n",
        "system.addForce(torchforce)\n",
        "\n",
        "\n",
        "# Create an integrator with a time step of 0.5 fs\n",
        "temperature = 298.15 * unit.kelvin\n",
        "frictionCoeff = 1 / unit.picosecond\n",
        "timeStep = 0.5 * unit.femtosecond\n",
        "integrator = mm.LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)\n",
        "\n",
        "# Create a simulation and set the initial positions and velocities\n",
        "simulation = app.Simulation(pdb.topology, system, integrator)\n",
        "simulation.context.setPositions(pdb.positions)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True, speed=True)\n",
        "simulation.reporters.append(reporter)\n",
        "\n",
        "simulation.reporters.append(app.PDBReporter('mace_traj.pdb', 100))\n",
        "\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Z4R3HYjRpiUl"
      },
      "source": [
        "## Extra exercises\n",
        "<a id=\"extra\"></a>\n",
        "\n",
        "- Use one of the other pretrained MACE models in the ANI_trained_MACE.zip file.\n",
        "- Make a mixed system using the MACE MLP\n",
        "- The openmm-ml createMixedSystem function has the (undocumented) ability to create a mixed system where the ML region can be interpolated by a lambda value between the ML and MM representation."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "dzLkCUJ0H-T0"
      },
      "source": [
        "## References\n",
        "<a id=\"references\"></a>\n",
        "\n",
        "[1] Towards chemical accuracy for alchemical free energy calculations with hybrid physics-based machine learning / molecular mechanics potentials,\n",
        "Dominic A. Rufa, Hannah E. Bruce Macdonald, Josh Fass, Marcus Wieder, Patrick B. Grinaway, Adrian E. Roitberg, Olexandr Isayev, John D. Chodera,\n",
        "bioRxiv 2020.07.29.227959; doi: https://doi.org/10.1101/2020.07.29.227959\n",
        "\n",
        "[2] NNP/MM: Fast molecular dynamics simulations with machine learning potentials and molecular mechanics,\n",
        "Raimondas Galvelis, Alejandro Varela-Rial, Stefan Doerr, Roberto Fino, Peter Eastman, Thomas E. Markland, John D. Chodera, Gianni De Fabritiis,\n",
        "arXiv:2201.08110; doi: https://doi.org/10.48550/arXiv.2201.08110\n"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Solutions"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "gpuType": "T4",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "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.10.12"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}