diff --git a/.github/workflows/tutorials.yml b/.github/workflows/tutorials.yml new file mode 100644 index 00000000..6231cda8 --- /dev/null +++ b/.github/workflows/tutorials.yml @@ -0,0 +1,41 @@ +name: "Test tutorials" + +on: + push: + branches: + - "master" + pull_request: + branches: + - "master" + schedule: + - cron: "0 0 * * 0" # Every Sunday at midnight + +jobs: + tutorial-1: + runs-on: "ubuntu-latest" + defaults: + run: + shell: bash -l {0} + steps: + - name: "Check out" + uses: actions/checkout@v2 + + - name: "Install dependencies with conda" + uses: conda-incubator/setup-miniconda@v2 + with: + activate-environment: test + environment-file: devtools/conda-envs/test-tutorials.yml + python-version: 3.7 # Mimics Google Colab + + - name: "List conda packages" + run: conda list + + - name: "Convert the tutorial" + run: | + jupyter nbconvert --to python --stdout tutorials/openmm-torch-nnpops.ipynb > test.py + sed -i '/condacolab/d' test.py # Remove conda-colab + sed -i 's#&> /dev/null##g' test.py # Remove the output redirection + cat test.py + + - name: "Run the tutorial" + run: ipython test.py \ No newline at end of file diff --git a/README.md b/README.md index 84c78b07..08e752c0 100644 --- a/README.md +++ b/README.md @@ -68,6 +68,11 @@ install the Python wrapper. Using the OpenMM PyTorch plugin =============================== +Tutorials +--------- + +- [A simple simulation of alanine dipeptide with ANI-2x using OpenMM-Torch and NNPOps](tutorials/openmm-torch-nnpops.ipynb) + Exporting a PyTorch model for use in OpenMM ------------------------------------------- diff --git a/devtools/conda-envs/test-tutorials.yml b/devtools/conda-envs/test-tutorials.yml new file mode 100644 index 00000000..d738d517 --- /dev/null +++ b/devtools/conda-envs/test-tutorials.yml @@ -0,0 +1,5 @@ +name: test +channels: + - conda-forge +dependencies: + - notebook \ No newline at end of file diff --git a/tutorials/openmm-torch-nnpops.ipynb b/tutorials/openmm-torch-nnpops.ipynb new file mode 100644 index 00000000..d479b744 --- /dev/null +++ b/tutorials/openmm-torch-nnpops.ipynb @@ -0,0 +1,440 @@ +{ + "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()" + ] + }, + { + "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: Remove \"mmh\" when NNPOps in available on conda-forge (https://github.com/openmm/NNPOps/issues/26)\n", + "# Note: \"cudatoollit=11.2\" is need to override a pinning in conda-colab\n", + "!conda install -q -c conda-forge -c mmh \\\n", + " openmm-torch nnpops torchani openmmtools \\\n", + " cudatoolkit=11.2 \\\n", + " &> /dev/null # Comment this line to see the installaion logs" + ] + }, + { + "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 +}