tutorials/openmm-torch-nnpops.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oRr7FSA13_Wv"
      },
      "source": [
        "# Tutorial: a simple simulation of alanine dipeptide with ANI-2x using OpenMM-Torch and NNPOps\n",
        "\n",
        "You can run this tutorial directly in your browser: [![Open On Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/openmm/openmm-torch/blob/master/tutorials/openmm-torch-nnpops.ipynb)\n",
        "\n",
        "Covered topics:\n",
        " * Installation of the software with [Conda](https://docs.conda.io/)\n",
        " * Creation of an NNP (neural network potential) with [TorchANI](https://aiqm.github.io/torchani/)\n",
        " * Acceleration of [ANI-2x](https://chemrxiv.org/engage/api-gateway/chemrxiv/assets/orp/resource/item/60c747e9567dfec574ec48df/original/extending-the-applicability-of-the-ani-deep-learning-molecular-potential-to-sulfur-and-halogens.pdf) with [NNPOps](https://github.com/openmm/NNPOps) \n",
        " * Integration of [OpenMM](https://openmm.org/) and [PyTorch](https://pytorch.org/) and with [OpenMM-Torch](https://github.com/openmm/openmm-torch)\n",
        " * Setup and execution of a simulation"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BKG1ZNq94f-Z"
      },
      "source": [
        "## Install Conda\n",
        "\n",
        "[Conda](https://docs.conda.io/) is a package and environment manager. On Google Colab, Conda is installed with [conda-colab](https://github.com/jaimergp/condacolab). On your computer, you should follow these [installation instructions](https://docs.conda.io/projects/conda/en/latest/user-guide/install/index.html).\n",
        "\n",
        "⚠️ Do not use conda-colab on your computer!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "EcB2lc1s3ZXG",
        "outputId": "1ce0bad8-a050-48d5-c18a-9e49fbeb11c2"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "⏬ Downloading https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh...\n",
            "📦 Installing...\n",
            "📌 Adjusting configuration...\n",
            "🩹 Patching environment...\n",
            "⏲ Done in 0:00:34\n",
            "🔁 Restarting kernel...\n"
          ]
        }
      ],
      "source": [
        "!pip install -q condacolab\n",
        "import condacolab\n",
        "condacolab.install_miniconda()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tWlPu4KF5DFZ"
      },
      "source": [
        "## Install software\n",
        "\n",
        "The [conda-forge](https://conda-forge.org/) channel is used for software.\n",
        "\n",
        "⚠️ The installation might take up to 10 min!"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "id": "-qFE_cKF5958"
      },
      "outputs": [],
      "source": [
        "# NOTE: PyTorch pinned because https://github.com/conda-forge/openmm-torch-feedstock/issues/20\n",
        "!mamba install -q -c conda-forge \\\n",
        "               openmm-torch nnpops torchani openmmtools \\\n",
        "               pytorch=1.11 \\\n",
        "               &> /dev/null # Comment this line to see a log\""
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZQwKdquUF5wo"
      },
      "source": [
        "## Prepare a simulation system\n",
        "\n",
        "For simplicity, the alanine dipeptide system from [OpenMM-Tools](https://openmmtools.readthedocs.io/en/latest/) is used."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "zNng_d4mGKve",
        "outputId": "218e0775-766b-4597-9d48-96e3d65df374"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "Warning: importing 'simtk.openmm' is deprecated.  Import 'openmm' instead.\n"
          ]
        }
      ],
      "source": [
        "import openmmtools\n",
        "\n",
        "# Get the system of alanine dipeptide\n",
        "ala2 = openmmtools.testsystems.AlanineDipeptideVacuum(constraints=None)\n",
        "\n",
        "# Remove MM forces\n",
        "while ala2.system.getNumForces() > 0:\n",
        "  ala2.system.removeForce(0)\n",
        "\n",
        "# The system should not contain any additional force and constrains\n",
        "assert ala2.system.getNumConstraints() == 0\n",
        "assert ala2.system.getNumForces() == 0\n",
        "\n",
        "# Get the list of atomic numbers\n",
        "atomic_numbers = [atom.element.atomic_number for atom in ala2.topology.atoms()]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BWlQTEF1IS4S"
      },
      "source": [
        "## Create a NNP\n",
        "\n",
        "A NNP (neural network potential) is created with [TorchANI](https://aiqm.github.io/torchani/). In this case [ANI-2x](https://chemrxiv.org/engage/api-gateway/chemrxiv/assets/orp/resource/item/60c747e9567dfec574ec48df/original/extending-the-applicability-of-the-ani-deep-learning-molecular-potential-to-sulfur-and-halogens.pdf) is used, which can be accelerated with [NNPOPs](https://github.com/openmm/NNPOps)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "ofemZo6sIbLU",
        "outputId": "4181fcb6-db72-47e7-edb5-2a4bb0252de7"
      },
      "outputs": [
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/site-packages/torchani/__init__.py:55: UserWarning: Dependency not satisfied, torchani.ase will not be available\n",
            "  warnings.warn(\"Dependency not satisfied, torchani.ase will not be available\")\n"
          ]
        },
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/site-packages/torchani/resources/\n",
            "Downloading ANI model parameters ...\n"
          ]
        },
        {
          "name": "stderr",
          "output_type": "stream",
          "text": [
            "/usr/local/lib/python3.7/site-packages/torch/functional.py:1069: UserWarning: torch.meshgrid: in an upcoming release, it will be required to pass the indexing argument. (Triggered internally at  /home/conda/feedstock_root/build_artifacts/pytorch-recipe_1640851336451/work/aten/src/ATen/native/TensorShape.cpp:2156.)\n",
            "  return _VF.cartesian_prod(tensors)  # type: ignore[attr-defined]\n"
          ]
        }
      ],
      "source": [
        "import torch as pt\n",
        "from torchani.models import ANI2x\n",
        "from NNPOps import OptimizedTorchANI\n",
        "\n",
        "class NNP(pt.nn.Module):\n",
        "\n",
        "  def __init__(self, atomic_numbers):\n",
        "\n",
        "    super().__init__()\n",
        "\n",
        "    # Store the atomic numbers\n",
        "    self.atomic_numbers = pt.tensor(atomic_numbers).unsqueeze(0)\n",
        "\n",
        "    # Create an ANI-2x model\n",
        "    self.model = ANI2x(periodic_table_index=True)\n",
        "\n",
        "    # Accelerate the model\n",
        "    self.model = OptimizedTorchANI(self.model, self.atomic_numbers)\n",
        "\n",
        "  def forward(self, positions):\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",
        "\n",
        "# Create an instance of the model\n",
        "nnp = NNP(atomic_numbers)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iOTneQ1_I9aB"
      },
      "source": [
        "At this point, the potential energy of the system can be computed with the NNP:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "D42I_L9r-Krj",
        "outputId": "c1738d48-71e2-40ea-9965-d9dbd2147406"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "-1301523.8702643516\n"
          ]
        }
      ],
      "source": [
        "# Comute the potential energy\n",
        "pos = pt.tensor(ala2.positions.tolist())\n",
        "energy_1 = float(nnp(pos))\n",
        "print(energy_1)\n",
        "\n",
        "# Check if the energy is correct\n",
        "assert pt.isclose(pt.tensor(energy_1), pt.tensor(-1301523.8703817206))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zOV53fN2sB3C"
      },
      "source": [
        "## Add the NNP to the system\n",
        "\n",
        "In order to use the NNP in a simulation, it has to loaded with [OpenMM-Torch](https://github.com/openmm/openmm-torch) and added to the system."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "id": "Snj9UrPbsIS5"
      },
      "outputs": [],
      "source": [
        "from openmmtorch import TorchForce\n",
        "\n",
        "# Save the NNP to a file and load it with OpenMM-Torch\n",
        "pt.jit.script(nnp).save('model.pt')\n",
        "force = TorchForce('model.pt')\n",
        "\n",
        "# Add the NNP to the system\n",
        "ala2.system.addForce(force)\n",
        "assert ala2.system.getNumForces() == 1"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "jLtSbKg6BB1n"
      },
      "source": [
        "## Setup a simulation\n",
        "\n",
        "Setup a simulation with [OpenMM](https://openmm.org/)."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "id": "6l2vPxPZ3u56"
      },
      "outputs": [],
      "source": [
        "import sys\n",
        "from openmm import LangevinMiddleIntegrator\n",
        "from openmm.app import Simulation, StateDataReporter\n",
        "from openmm.unit import kelvin, picosecond, femtosecond\n",
        "\n",
        "# Create an integrator with a time step of 1 fs\n",
        "temperature = 298.15 * kelvin\n",
        "frictionCoeff = 1 / picosecond\n",
        "timeStep = 1 * femtosecond\n",
        "integrator = LangevinMiddleIntegrator(temperature, frictionCoeff, timeStep)\n",
        "\n",
        "# Create a simulation and set the initial positions and velocities\n",
        "simulation = Simulation(ala2.topology, ala2.system, integrator)\n",
        "simulation.context.setPositions(ala2.positions)\n",
        "# simulation.context.setVelocitiesToTemperature(temperature) # This does not work (https://github.com/openmm/openmm-torch/issues/61)\n",
        "\n",
        "# Configure a reporter to print to the console every 0.1 ps (100 steps)\n",
        "reporter = StateDataReporter(file=sys.stdout, reportInterval=100, step=True, time=True, potentialEnergy=True, temperature=True)\n",
        "simulation.reporters.append(reporter)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "4ProIzIWL4Co"
      },
      "source": [
        "At this point, the potential energy of the system can be computed again:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "JN9v2rz3KbTz",
        "outputId": "15367844-1caa-40c1-bb5b-f94f7b51eb1b"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "-1301523.8702643516\n"
          ]
        }
      ],
      "source": [
        "from openmm.unit import kilojoule_per_mole\n",
        "\n",
        "# Comute the potential energy\n",
        "state = simulation.context.getState(getEnergy=True)\n",
        "energy_2 = state.getPotentialEnergy().value_in_unit(kilojoule_per_mole)\n",
        "print(energy_2)\n",
        "\n",
        "# Check if the energy is correct\n",
        "assert pt.isclose(pt.tensor(energy_1), pt.tensor(energy_2))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "3DwdnGQ8NT2X"
      },
      "source": [
        "## Run the simulation\n",
        "\n",
        "Run your first NNP simulation.\n",
        "\n",
        "⚠️ The simulations are not deterministic! Each time the log will be a bit different."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 7,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "da0v6qJHBtUb",
        "outputId": "e477991e-93b0-4e72-de4f-d7a8b05052de"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "#\"Step\",\"Time (ps)\",\"Potential Energy (kJ/mole)\",\"Temperature (K)\"\n",
            "100,0.10000000000000007,-1301527.998092823,62.28599411217165\n",
            "200,0.20000000000000015,-1301519.2141580286,68.6151680329106\n",
            "300,0.3000000000000002,-1301513.905558333,103.32518127627559\n",
            "400,0.4000000000000003,-1301517.7584694924,118.20611762387\n",
            "500,0.5000000000000003,-1301514.9794064017,165.4939713543869\n",
            "600,0.6000000000000004,-1301507.6431399288,148.54034468306415\n",
            "700,0.7000000000000005,-1301511.109281123,157.43310789810866\n",
            "800,0.8000000000000006,-1301498.8263809385,157.8375886414082\n",
            "900,0.9000000000000007,-1301511.3014532926,199.2036328097003\n",
            "1000,1.0000000000000007,-1301485.2554342675,148.26160451888617\n"
          ]
        }
      ],
      "source": [
        "# Run the simulations for 1 ps (1000 steps)\n",
        "simulation.step(1000)"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [
        "BKG1ZNq94f-Z"
      ],
      "name": "openmm-torch-nnpops.ipynb",
      "provenance": []
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python",
      "version": "3.7.12"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}