section_3/machine_learning_potentials.ipynb

{
  "cells": [
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "HL9clVbBmMu-"
      },
      "source": [
        "# Machine Learning Potentials\n",
        "\n",
        "You can run this notebook in your browser: \n",
        "\n",
        "[![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openmm/openmm_workshops/blob/main/section_3/machine_learning_potentials.ipynb)\n",
        "\n",
        "## Table of contents\n",
        "- Introduction\n",
        "- OpenMM ML software\n",
        "- Basics of OpenMM compatible MLPs\n",
        "- Installing\n",
        "- Exporting a PyTorch model for use in OpenMM\n",
        "- Simulation of alanine dipeptide with ANI-2x using OpenMM-Torch\n",
        "    - Prepate test system\n",
        "    - Create a NNP\n",
        "    - Adding the NNP to an OpenMM simulation\n",
        "- Mixed MM/ML system\n",
        "    - Creating the system\n",
        "    - create the MLP\n",
        "- Using OpenMM-ML package\n",
        "- Using NNPOps\n",
        "- Implementing other models - MACE\n",
        "- Extra exercises\n",
        "- References\n",
        "- 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)[3], Gaussian process regression (GAP)[4], graph neural network (MACE)[5], equivariant transformer (TorchMD-NET)[6], or something else. MLPs are trained on first-principles quantum mechanical (QM) methods such as DFT, and they can be used in atomistic MD simulations in the place of a classical force fields, bringing the accuracy of QM methods without the computational expense."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## OpenMM Machine Learning Software\n",
        "<a id=\"openmmmlsoftware\"></a>\n",
        "\n",
        "OpenMM is an MD engine, so we will focus  be covering how to run simulations using MLPs, and not 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, i.e., an [OpenMM Force class](http://docs.openmm.org/latest/api-python/library.html#forces) that computes a contribution to the total 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 commonly used in popular neural network potentials that can be used to speed up your MLP implementation.\n",
        "\n",
        "If you were to use your own model, you would first use `openmm-torch` to interface your MLP with OpenMM. You could then try and optimize it using operations from `NNPOps`. When you are ready 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 reads a set of atomic coordinates as input and outputs the potential energy. The forces on each particle are computed by backpropagation, 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 that can be exported to [TorchScript]([https://pytorch.org/docs/stable/jit.html#torchscript]). The model takes a $(N,3)$-shaped tensor of particle positions, where $N$ is the number of particles, and outputs the potential energy. `openmm-torch` calculates the forces using [Autograd](https://pytorch.org/docs/stable/autograd.html). Optionally, the model can return the forces, which `openmm-torch` will then use."
      ]
    },
    {
      "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`NNPOps`.\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "  ⚠️ <b>Due to the added complexity of creating a conda environment with PyTorch, we recommend you run this notebook in Colab. It should work on Linux and MacOS though. This tutorial will not work on Windows!</b>\n",
        "</div>\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",
        "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!')"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-block alert-info\">\n",
        "  ⚠️ <b>Note: During this step in Colab, the kernel will restart, which may trigger the error message: \"Your session crashed for an unknown reason.\" This is expected behavior and can be safely ignored.</b>\n",
        "</div>\n",
        "\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "⚠️ <b>Note that the installation will take several minutes!</b>\n",
        "</div>"
      ]
    },
    {
      "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",
        "    %env CONDA_OVERRIDE_CUDA=11.8\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"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Download the files needed for this tutorial."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "!wget https://raw.githubusercontent.com/openmm/openmm_workshops/main/section_3/alanine-dipeptide.pdb\n",
        "!wget https://raw.githubusercontent.com/openmm/openmm_workshops/main/section_3/workshop_utils.py"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## PyTorch\n",
        "\n",
        "PyTorch is an open-source machine learning framework that is primarily used for developing and training deep learning models. It provides a dynamic computational graph that allows users to define and modify neural networks on the fly, making it flexible and efficient for building complex models.\n",
        "\n",
        "Some key features of PyTorch include:\n",
        "\n",
        " - Dynamic computational graph: PyTorch uses a tape-based automatic differentiation system, which enables users to define and modify models on-the-fly. This dynamic nature makes it easy to debug and experiment with different model architectures.\n",
        "\n",
        " - GPU acceleration: PyTorch leverages the power of graphics processing units (GPUs) to accelerate the training and inference processes. It provides seamless integration with CUDA, a parallel computing platform, allowing for efficient computation on GPUs.\n",
        "\n",
        " - Neural network modules: PyTorch provides a rich library of pre-defined modules and functions for building neural networks. These modules include layers, activation functions, loss functions, and optimization algorithms, making it easy to construct and train deep learning models.\n",
        "\n",
        " - Community and ecosystem: PyTorch has a vibrant community of developers and researchers who contribute to its development and share their work. This community has created numerous resources, such as tutorials, libraries, and pre-trained models, which can be readily used for various machine learning tasks.\n",
        "\n",
        " - Deployment options: PyTorch offers various deployment options, including exporting models for inference in production environments. It provides tools like TorchScript, which allows models to be serialized and executed independently of the Python runtime, enabling deployment on platforms with limited resources.\n",
        "\n",
        "PyTorch has gained popularity due to its ease of use, flexibility, and extensive support for research and development in the field of deep learning. It is widely used by researchers, academics, and industry professionals for a range of applications, including computer vision, natural language processing, reinforcement learning, and our use case of atomic Force Fields. \n",
        "\n",
        "We will be using pre-trained models so we do not need to know much about PyTorch. We just need to know that you can create models that take in atomic coordinates and predict the energy. It you want to learn more about PyTorch there are plenty of tutorials available: https://pytorch.org/tutorials/index.html"
      ]
    },
    {
      "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 as `torch.Tensor` of shape `(N, 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 simple 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 (N,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 write the code in the `forward` method. The forward method defines the computation performed at every call of the model.\n",
        "\n",
        "Now that we have defined a model, we can create an instance of it."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "force_module = ForceModule()"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can give it some example input"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Create random tensor of 4 coordinates\n",
        "# We specify we want to record the gradient so we can compute the forces\n",
        "input = torch.rand((4,3), requires_grad=True)\n",
        "print(\"Input:\", input)\n",
        "energy = force_module(input)\n",
        "print(\"Energy:\", energy)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can calculate the forces resulting from this harmonic potential using Autograd."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "# Backward pass computes the gradients\n",
        "energy.backward()\n",
        "# We can access the gradients with the grad attribute\n",
        "# F = - grad (Potential)\n",
        "forces = -input.grad \n",
        "print(\"Forces:\", forces)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-LYilcYPmMvE"
      },
      "source": [
        "To export the model for use in OpenMM (or other software), it must be converted to a TorchScript module (and optionally 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",
        "scripted_module = torch.jit.script(force_module)\n",
        "print(scripted_module)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can then save the module to a file. The saving process serializes the module, allowing us to load it at any time into the C++ API as required by `openmm-torch`'s `TorchForce`. For more information, see the [PyTorch documentation](https://pytorch.org/docs/stable/generated/torch.jit.save.html)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lK91t53WmMvF"
      },
      "outputs": [],
      "source": [
        "# Save the serialized compute graph to a file\n",
        "scripted_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",
        "print(\"number of forces = \", system.getNumForces())\n",
        "\n",
        "# Add the TorchForce to the system\n",
        "system.addForce(torch_force)\n",
        "print(\"number of forces = \", system.getNumForces())"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "-Ym2SuXPmMvG"
      },
      "source": [
        "Now that we know how to use a PyTorch model in an OpenMM simulation, we are ready to learn how to create a neural network potential (NNP) that uses ANI-2x."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Simulation of Alanine Dipeptide with ANI-2x Using OpenMM-Torch\n",
        "<a id=\"ani\"></a>\n",
        "\n",
        "ANI-2x is a general NNP that works with molecules containing H, C, N, O, F, Cl, and S. For more information, please refer to the publication [3]. The model is available in 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 constraints\n",
        "assert system.getNumConstraints() == 0\n",
        "assert system.getNumForces() == 0\n",
        "\n",
        "# Get the list of atomic numbers. \n",
        "# We will use this list 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 are now ready to 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",
        "    \"\"\"A simple wrapper around the ANI-2x model from torchani.\"\"\"\n",
        "\n",
        "    def __init__(self, atomic_numbers: torch.Tensor):\n",
        "        \"\"\"\n",
        "        Initialize the NNP model.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        atomic_numbers : torch.Tensor with shape (N,)\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",
        "        \"\"\"The forward method returns the energy computed from positions.\n",
        "\n",
        "        Parameters\n",
        "        ----------\n",
        "        positions : torch.Tensor with shape (N,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",
        "        # Convert the atomic positions to Angstrom\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"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The `NNP` looks rather complex so we will break it down line by line.\n",
        "\n",
        "The first line\n",
        "```python\n",
        "class NNP(torch.nn.Module):\n",
        "```\n",
        "defines the NNP as a python [`Class`](https://docs.python.org/3/tutorial/classes.html) called `NNP` that inherits from the [`torch.nn.Module`](https://pytorch.org/docs/stable/generated/torch.nn.Module.html) class.\n",
        "\n",
        "```python\n",
        "def __init__(self, atomic_numbers: torch.Tensor):\n",
        "```\n",
        "Is the constructor definition. It states that when we create a new instance of the class, we must pass a `torch.Tensor` of the atomic numbers of the system as an argument. All the code within the constructor is called when the NNP is first created.\n",
        "\n",
        "```python\n",
        "super().__init__()\n",
        "```\n",
        "Calls the constructor of the parent class, viz. `torch.nn.Module`.\n",
        "\n",
        "```python\n",
        "# Store the atomic numbers\n",
        "self.atomic_numbers = atomic_numbers.unsqueeze(0)\n",
        "```\n",
        "Stores the provided `atomic_numbers` tensor as an attribute of the class. The `unsqueeze(0)` converts the tensor from 1D with size $(N)$ to 2D with size $(1,N)$. This is know as adding a batch dimension. This needs to be done because the model we will use expects batched data (even if the batch size is 1).\n",
        "\n",
        "```python\n",
        "# Create an ANI-2x model\n",
        "self.model = ANI2x(periodic_table_index=True)\n",
        "```\n",
        "This creates an instance of an ANI2x model from the [TorchANI](https://aiqm.github.io/torchani/api.html#module-torchani.models) package.\n",
        "\n",
        "```python\n",
        "# Make sure it is on the same device at the atomic_numbers tensor\n",
        "self.model.to(self.atomic_numbers.device)\n",
        "```\n",
        "\n",
        "This makes sure the model is on the same device as the `atomic_numbers` tensor. The [device](https://pytorch.org/docs/stable/tensor_attributes.html#torch.device) is either 'cpu' or 'cuda' (GPU). \n",
        "\n",
        "```python\n",
        "def forward(self, positions: torch.Tensor):\n",
        "```\n",
        "This defines the [forward method](https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.forward). The forward method is the code that is called every time the model is evaluated. We define that a `torch.Tensor` of the atomic positions must be passed as an argument when we evaluate the model.\n",
        "\n",
        "```python\n",
        "# Prepare the positions\n",
        "positions = positions.unsqueeze(0).float() * 10 # nm --> Å\n",
        "```\n",
        "This adds a batch dimension to the positions tensor, converts it to floating point precision (it might be in double if OpenMM is running in double precision), and converts the units from the OpenMM default of nanometers to Angstrom, as required by the ANI model.\n",
        "\n",
        "```python\n",
        "# Run ANI-2x\n",
        "result = self.model((self.atomic_numbers, positions))\n",
        "\n",
        "energy = result.energies[0] * 2625.5 # Hartree --> kJ/mol\n",
        "```\n",
        "This evaluates the ANI-2x model on the atomic configuration. The result will contain the total potential energy. We need to convert from the ANI units of Hartree to the OpenMM units of kJ/mol.\n",
        "\n",
        "Implementing other NNPs will follow a similar format. The key point is that you must convert between the OpenMM format and the format expected by the model."
      ]
    },
    {
      "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. If we print the model, we can see the underlying neural network architecture."
      ]
    },
    {
      "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",
        "print(nnp)"
      ]
    },
    {
      "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 compile the model as TorchScript code 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_module = torch.jit.script(nnp)\n",
        "torchforce = TorchForce(torch_module)\n",
        "\n",
        "# Add the NNP to the system\n",
        "system.addForce(torchforce)\n",
        "\n",
        "print(\"number of forces = \", system.getNumForces())\n",
        "assert(system.getNumForces() == 1)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-block alert-info\">\n",
        "⚠️ <b>There should be only 1 force in the system. If the assertion above fails, it might be because you have run the cell multiple times. You must go back and run the \"Prepare a Test System\" cell to create the system with no forces!</b>\n",
        "</div>"
      ]
    },
    {
      "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",
        "simulation.context.setVelocitiesToTemperature(temperature)\n",
        "\n",
        "# Add a PDB reporter to save the trajectory to a file\n",
        "simulation.reporters.append(app.PDBReporter('traj.pdb', 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)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Z9cJiI0xmMvH"
      },
      "source": [
        "Now, we can compute the energy and forces again using the OpenMM interface and compare them to the energy and forces computed with the PyTorch interface. This is a useful check, as bugs can sometimes arise during the export, serialization, and loading steps."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "CkFPCAiUmMvH",
        "outputId": "364ed990-da27-4a65-ae01-0d012c235bb7"
      },
      "outputs": [],
      "source": [
        "import numpy as np\n",
        "\n",
        "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",
        "\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": {},
      "source": [
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 1</b>\n",
        "\n",
        "Download the `'traj.pdb'` file and visualize it. You should see alanine dipeptide moving around. \n",
        "</div>"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "ye5bo0PVavZI"
      },
      "source": [
        "## Mixed System\n",
        "<a id=\"mixedsystem\"></a>\n",
        "\n",
        "In general, ML force fields are still too computationally expensive 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 forcefield. This approach is similar to hybrid QM/MM methods.\n",
        "\n",
        "We will not learn how to create hybrid ML/MM systems in OpenMM. Our example system will be the alanine-dipeptide in a water box.\n",
        "The intramolecular forces of the alanine dipeptide will be modeled with the ANI-2x MLP, while the water molecules and the intermolecular interactions between the alanine dipeptide and the water will be modeled with an MM force field.\n",
        "\n",
        "The total potential energy of the mixed system can be written as\n",
        "\n",
        "$V_{MM/ML}(r) = V_{MM}(r_{MM}) + V_{MM-ML}(r) + V_{ML}(r_{ML})$ \n",
        "\n",
        "where $r$ are the coordinates of all atoms, $r_{MM}$ are the coordinates of the atoms of the MM region, and $r_{ML}$ are the coordinates of the atoms of 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.\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "⚠️ <b>The mixed system strategy we have outlined and will implement is only appropriate when the ML region is a whole single molecule. ML/MM divisions spliting the system along a chemical bond are not yet supported.</b>\n",
        "</div>"
      ]
    },
    {
      "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 workshop_utils import removeBonds\n",
        "\n",
        "def removeMMInteraction(system, ml_atoms):\n",
        "    \"\"\"\n",
        "    Remove the MM interactions between the ML atoms.\n",
        "\n",
        "    Parameters\n",
        "    ----------\n",
        "    system : openmm.System\n",
        "        The system object\n",
        "    ml_atoms : list of int\n",
        "        The list of atom indices in the ML region\n",
        "\n",
        "    Returns\n",
        "    -------\n",
        "    newSystem : openmm.System\n",
        "        The new system with the MM interactions between the ML atoms removed\n",
        "    \"\"\"\n",
        "    # Remove the bonded interactions within 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(atomList[i], atomList[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",
        "                a1 = atomList[i]\n",
        "                for j in range(i):\n",
        "                    a2 = atomList[j]\n",
        "                    if (a1, a2) not in existing and (a2, a1) not in existing:\n",
        "                        force.addExclusion(a1, a2, True)\n",
        "\n",
        "    return newSystem"
      ]
    },
    {
      "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.\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 2</b>\n",
        "\n",
        "Add the code that stores the atomic numbers and the code that creates the ANI-2x model. \n",
        "\n",
        "Tip: this part is the same as the previous example.\n",
        "</div>"
      ]
    },
    {
      "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",
        "    def __init__(self, atomic_numbers: torch.Tensor, ml_atoms: torch.Tensor):\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",
        "        FIXME\n",
        "\n",
        "        # Create an ANI-2x model\n",
        "        FIXME\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",
        "        # 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"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "CkdqLIvCPRAL"
      },
      "source": [
        "Now we can create an instance of the MLP and add it to the system.\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 3</b>\n",
        "\n",
        "Add the `TorchForce` to the system.\n",
        "</div>"
      ]
    },
    {
      "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",
        "# Compile the NNP to TorchScript and load it with OpenMM-Torch\n",
        "torch_module = torch.jit.script(hybridnnp)\n",
        "\n",
        "from openmmtorch import TorchForce\n",
        "torchforce = TorchForce(torch_module)\n",
        "\n",
        "# Make the mixed system\n",
        "mixed_system = removeMMInteraction(system, ml_atoms.tolist())\n",
        "\n",
        "# Add the TorchForce\n",
        "FIXME\n",
        "\n",
        "# Print out the forces\n",
        "for force in mixed_system.getForces():\n",
        "    print(force)\n",
        "\n",
        "assert(mixed_system.getNumForces()==6)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "taMDOGCpUaqd"
      },
      "source": [
        "### Simulate the Hybrid ML/MM System"
      ]
    },
    {
      "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",
        "simulation.context.setVelocitiesToTemperature(temperature)\n",
        "\n",
        "# Minimize the system\n",
        "simulation.minimizeEnergy(maxIterations=100)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps and write to a PDB file\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)\n",
        "simulation.reporters.append(app.PDBReporter('mixed_traj.pdb', 100, enforcePeriodicBox=False))\n",
        "\n",
        "# Simulate 0.5 ps\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 4</b>\n",
        "\n",
        "Visualize `'mixed_traj.pdb'`.\n",
        "</div>"
      ]
    },
    {
      "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 an MLP to a system. As you’ve seen, it involves quite a lot of code. The [openmm-ml](https://github.com/openmm/openmm-ml) package was created as a high-level interface to simplify the use of pre-trained ML models in OpenMM simulations. We will now run 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 the alanine dipeptide molecule and simulate it in vacuum with ANI-2x with the `MLPotential.createSystem` function from the `openmm-ml` package."
      ]
    },
    {
      "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",
        "print(pdb.topology)\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",
        "# 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",
        "simulation.context.setVelocitiesToTemperature(temperature)\n",
        "\n",
        "# Minimize the system\n",
        "simulation.minimizeEnergy(maxIterations=100)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps and write to a PDB file\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)\n",
        "simulation.reporters.append(app.PDBReporter('traj.pdb', 100))\n",
        "\n",
        "# Simulate 0.5 ps\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "<div class=\"alert alert-block alert-info\">\n",
        "⚠️ <b>You can safely ignore the error message that says `failed to equip 'nnpops' with error: No module named 'NNPOps'`.</b>\n",
        "</div>"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "0VIiA4tSeTBJ"
      },
      "source": [
        "### Create a Mixed System\n",
        "We can just as easily create a mixed ML/MM system.\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 5</b>\n",
        "\n",
        "Write the code to create the potential using MLPotential. \n",
        "\n",
        "Tip: Look at the OpenMM-ML readme: https://github.com/openmm/openmm-ml\n",
        "</div>"
      ]
    },
    {
      "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",
        "import sys\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",
        "# 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",
        "FIXME\n",
        "\n",
        "# Create the mixed system\n",
        "chains = list(modeller.topology.chains())\n",
        "ml_atoms = [atom.index for atom in chains[0].atoms()]\n",
        "mixed_system = potential.createMixedSystem(modeller.topology, mm_system, ml_atoms)\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.context.setVelocitiesToTemperature(temperature)\n",
        "\n",
        "# Minimize the system\n",
        "simulation.minimizeEnergy(maxIterations=100)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps and write to a PDB file\n",
        "reporter = app.StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)\n",
        "simulation.reporters.append(app.PDBReporter('traj.pdb', 100, enforcePeriodicBox=False))\n",
        "\n",
        "# Simulate 0.5 ps\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. NNPOps can be used to speed up ANI simulations. We can install it from conda-forge and use it through the `openmm-ml` interface.\n",
        "\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "⚠️ <b>NNPOps is not available on Windows.</b>\n",
        "</div>"
      ]
    },
    {
      "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 and use the nnpops implementation\n",
        "potential = MLPotential('ani2x', implementation='nnpops')\n",
        "\n",
        "# create a system that uses the MLP\n",
        "system = potential.createSystem(pdb.topology)\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",
        "simulation.context.setVelocitiesToTemperature(temperature)\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",
        "# Simulate 0.5 ps\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "zBe8H9m5hP6A"
      },
      "source": [
        "The `implementation='nnpops'` argument tells `MLPotential` to replace certain TorchANI PyTorch functions with their optimized NNPOps counterparts. If you are interesting in seeing how this is implemented, checkout out the [GitHub repository](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`. In this tutorial, we will show how to implement a MACE model.\n",
        "\n",
        "### Install software\n",
        "MACE can be installed via pip."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "txm22xKliRRl"
      },
      "outputs": [],
      "source": [
        "!pip install mace-torch"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Ong6DOIzia1r"
      },
      "source": [
        "### Get a Pretrained Model\n",
        "\n",
        "We are going to use the small pre-trained model from the MACE-OFF23 family, a transferable force field for organic molecules [7]. You can find the remaining versions [here](https://github.com/ACEsuit/mace-off/tree/main/mace_off23)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "anI4ZqDyjD5J",
        "outputId": "d74bdd41-f3bd-4152-fc7c-ee050aafae81"
      },
      "outputs": [],
      "source": [
        "# Download the model\n",
        "!wget https://raw.githubusercontent.com/ACEsuit/mace-off/main/mace_off23/MACE-OFF23_small.model"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "EE9DpkY7jmtK"
      },
      "source": [
        "### Define the MLP\n",
        "\n",
        "We will create a MACE MLP class, similar to what we did previously for ANI-2x. The MACE model, however, requires additional code to convert atomic numbers and positions into the necessary format. Additionally, the `simple_nl` function for calculating the neighbor list is available in the `workshop_utils` module."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "7Y2EBX4TkNdl"
      },
      "outputs": [],
      "source": [
        "import torch\n",
        "from workshop_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: \n",
        "            # Use cuda if available\n",
        "            self.device=torch.device(\"cuda\") if torch.cuda.is_available() else torch.device(\"cpu\")\n",
        "        else: \n",
        "            # Unless user has specified the device\n",
        "            self.device=torch.device(device)\n",
        "\n",
        "        self.default_dtype = dtype\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.energy_to_kJ = 96.49  # eV->kJ\n",
        "\n",
        "        # Load the model\n",
        "        self.model = torch.load(model_path, map_location=device)\n",
        "        \n",
        "        # Extract model parameters, define atomic number table, and compile the model\n",
        "        self.r_max = self.model.r_max\n",
        "        self.z_table = utils.AtomicNumberTable([int(z) for z in self.model.atomic_numbers])\n",
        "        self.model = jit.compile(self.model).to(self.default_dtype)\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",
        "            ).to(self.default_dtype)\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",
        "    def forward(self, positions, boxvectors: Optional[torch.Tensor] = None):\n",
        "        # Setup positions\n",
        "        positions = positions.to(device=self.device, dtype=self.default_dtype)\n",
        "        if self.indices is not None:\n",
        "            positions = positions[self.indices]\n",
        "        positions = positions*self.nm_to_distance\n",
        "\n",
        "        # Setup cell and pbc\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",
        "        # Calculate the shifts and edge_index\n",
        "        mapping, shifts_idx = simple_nl(positions, cell, pbc, self.r_max)\n",
        "        edge_index = torch.stack((mapping[0], mapping[1])).to(torch.int64)\n",
        "        shifts = torch.mm(shifts_idx, cell).to(self.default_dtype)\n",
        "\n",
        "        # Create input dict\n",
        "        input_dict = {\n",
        "            \"ptr\" : self.ptr,\n",
        "            \"node_attrs\": self.node_attrs,\n",
        "            \"batch\": self.batch,\n",
        "            \"pbc\": self.pbc,\n",
        "            \"positions\": positions,\n",
        "            \"edge_index\": edge_index,\n",
        "            \"shifts\": shifts,\n",
        "            \"cell\": cell,\n",
        "        }\n",
        "\n",
        "        # Predict the energy\n",
        "        energy = self.model(input_dict, compute_force=False)[\"interaction_energy\"]\n",
        "\n",
        "        assert energy is not None, \"The model did not return any energy. Please check the input.\"\n",
        "\n",
        "        # Return energy\n",
        "        energy = energy*self.energy_to_kJ\n",
        "\n",
        "        return energy"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "IKKPko8tmrwq"
      },
      "source": [
        "### Use the MACE MLP\n",
        "\n",
        "The rest of the code is then similar to using 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",
        "\n",
        "# Remove MM forces\n",
        "while system.getNumForces() > 0:\n",
        "  system.removeForce(0)\n",
        "\n",
        "# The system should not contain any additional force and constraints\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 = \"MACE-OFF23_small.model\"\n",
        "pbc = False \n",
        "indices = None # None means all atoms are used\n",
        "device = 'cuda' if torch.cuda.is_available() else 'cpu'\n",
        "\n",
        "# Create the MACE MLP\n",
        "mace_mlp = MACEForce(model_path, atomic_numbers, indices, pbc, device)\n",
        "\n",
        "# Compile the NNP to TorchScript and load it with OpenMM-Torch\n",
        "torch_module = torch.jit.script(mace_mlp)\n",
        "torchforce = TorchForce(torch_module)\n",
        "\n",
        "# Add it to the system\n",
        "system.addForce(torchforce)\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",
        "simulation.context.setVelocitiesToTemperature(temperature)\n",
        "\n",
        "# Configure a reporter to print to the console every 100 steps and write to a PDB file\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",
        "simulation.reporters.append(app.PDBReporter('mace_traj.pdb', 100))\n",
        "\n",
        "# Simulate 0.5 ps\n",
        "simulation.step(1000)"
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {
        "id": "Z4R3HYjRpiUl"
      },
      "source": [
        "## Extra exercises\n",
        "<a id=\"extra\"></a>\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 6</b>\n",
        "\n",
        "Use one of the other pretrained MACE-OFF23 models.\n",
        "\n",
        "Tip: For example, to download the medium-sized model, run `!wget https://raw.githubusercontent.com/ACEsuit/mace-off/main/mace_off23/MACE-OFF23_medium.model`.\n",
        "</div>\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 7</b>\n",
        "\n",
        "Increase the size of the water box in the mixed system. Measure the performance (ns/day) to find out how many MM atoms there need to be before the speed of the MM part becomes significant compared to the speed of the ML part.\n",
        "</div>\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 8 (Hard)</b>\n",
        "\n",
        "Make a mixed system using the MACE MLP.\n",
        "</div>\n",
        "<div class=\"alert alert-block alert-info\">\n",
        "ℹ️ <b>Exercise 9 (Hard)</b>\n",
        "\n",
        "The `openmm-ml` `createMixedSystem` function has the ability to create a mixed system where the ML region can be interpolated by a lambda value between the ML and MM representations. Look at the [API documentation](https://github.com/openmm/openmm-ml/blob/d5120bd1fe8cd7330bb3169f3549fd2d550d4c39/openmmml/mlpotential.py#L181) in the source code and try to use this functionality. \n",
        "</div>"
      ]
    },
    {
      "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",
        "\n",
        "[3] Xiang Gao, Farhad Ramezanghorbani, Olexandr Isayev, Justin S. Smith, and Adrian E. Roitberg, Chem. Inf. Model. 60, 7, 3408–3415 (2020), https://doi.org/10.1021/acs.jcim.0c00451 | https://aiqm.github.io/torchani/\n",
        "\n",
        "[4] AP Bartók, MC Payne, R Kondor, G Csányi, Physical review letters 104 (13), 136403 (2010), https://link.aps.org/doi/10.1103/PhysRevLett.104.136403\n",
        "\n",
        "[5] I. Batatia, D. P. Kovacs, G. Simm, C. Ortner, and G. Csányi. Advances in Neural Information \n",
        "    Processing Systems 35, 11423 (2022). https://github.com/ACEsuit/mace\n",
        "\n",
        "[6] P Thölke, G De Fabritiis, International Conference on Learning Representations, 2021, https://doi.org/10.48550/arXiv.2202.02541\n",
        "\n",
        "[7] Kovács, D. P.; Moore, J. H.; Browning, N. J.; Batatia, I.; Horton, J. T.; Kapil, V.; Witt, W. C.; Magdău, I.-B.; Cole, D. J.; Csányi, G. MACE-OFF23: Transferable Machine Learning Force Fields for Organic Molecules. arXiv December 29, 2023. https://doi.org/10.48550/arXiv.2312.15211."
      ]
    },
    {
      "attachments": {},
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Solutions\n",
        "\n",
        "*Exercise 2*\n",
        "```python\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",
        "\n",
        "*Exercise 3*\n",
        "```python\n",
        "mixed_system.addForce(torchforce)\n",
        "```\n",
        "\n",
        "*Exercise 5*\n",
        "```python\n",
        "# Create the MLP using ANI-2x\n",
        "potential = MLPotential('ani2x')\n",
        "```"
      ]
    }
  ],
  "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.9.16"
    },
    "orig_nbformat": 4
  },
  "nbformat": 4,
  "nbformat_minor": 0
}