OpenMM in Google Colab

Colab_OpenMM.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": [],
      "gpuType": "T4"
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "language_info": {
      "name": "python"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "source": [
        "# OpenMM in Google Colab"
      ],
      "metadata": {
        "id": "79hesOIpe5LY"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "%%time\n",
        "#@title Setup OpenMM\n",
        "import sys, os\n",
        "from sys import version_info\n",
        "PYTHON_VERSION = f\"{version_info.major}.{version_info.minor}\"\n",
        "\n",
        "if not os.path.isfile(\"CONDA_READY\"):\n",
        "  print(\"Installing Conda...\")\n",
        "  os.system(\"wget -qnc https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh\")\n",
        "  os.system(\"bash Miniforge3-Linux-x86_64.sh -bfp /usr/local\")\n",
        "  os.system(\"mamba config --set auto_update_conda false\")\n",
        "  os.system(\"touch CONDA_READY\")\n",
        "\n",
        "if not os.path.isfile(\"OPENMM_READY\"):\n",
        "  print(\"Installing OpenMM...\")\n",
        "  os.system(f\"mamba install -y -q -c omnia openmm python='{PYTHON_VERSION}' 2>&1 1>/dev/null\")\n",
        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/openmm openmm\")\n",
        "  os.system(\"touch OPENMM_READY\")\n",
        "\n",
        "PATH = f\"/usr/local/lib/python{PYTHON_VERSION}/site-packages/\"\n",
        "if PATH not in sys.path:\n",
        "  sys.path.insert(0, PATH)"
      ],
      "metadata": {
        "cellView": "form",
        "id": "-7mYTC6Gx9MP"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "from openmm import *\n",
        "from openmm.app import *\n",
        "from openmm.unit import *\n",
        "\n",
        "from sys import stdout\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt"
      ],
      "metadata": {
        "id": "09En8MZ2izIn"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Set up platform\n",
        "print(\"Setting up the platform...\")\n",
        "platform = Platform.getPlatformByName('CUDA')\n",
        "platform.setPropertyDefaultValue('Precision', 'mixed')"
      ],
      "metadata": {
        "id": "SWLT3z3te_A7"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Load PDB\n",
        "os.system(\"wget https://files.rcsb.org/view/2VIR.pdb\")\n",
        "\n",
        "structure_path = \"2VIR.pdb\"\n",
        "\n",
        "if structure_path.endswith(\".pdb\"):\n",
        "  pdb = PDBFile(structure_path)\n",
        "\n",
        "if structure_path.endswith(\".cif\"):\n",
        "  pdb = PDBxFile(structure_path)"
      ],
      "metadata": {
        "id": "YQTevTUJe45i"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Specify the forcefield\n",
        "forcefield = ForceField('amber14-all.xml', 'amber14/tip3pfb.xml')"
      ],
      "metadata": {
        "id": "Zz3cmdqSe--W"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Initialize Modeller\n",
        "modeller = Modeller(pdb.topology, pdb.positions)\n",
        "\n",
        "## Reset Water\n",
        "modeller.deleteWater()\n",
        "residues = modeller.addHydrogens(forcefield)\n",
        "modeller.addSolvent(forcefield, padding=1.0*nanometer)"
      ],
      "metadata": {
        "id": "f4Kabwd1e-78"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Create the system\n",
        "system = forcefield.createSystem(modeller.topology, nonbondedMethod=PME, nonbondedCutoff=1.0*nanometer, constraints=HBonds)\n",
        "integrator = LangevinMiddleIntegrator(300*kelvin, 1/picosecond, 0.004*picoseconds)\n",
        "simulation = Simulation(modeller.topology, system, integrator, platform)\n",
        "simulation.context.setPositions(modeller.positions)"
      ],
      "metadata": {
        "id": "aFztXAY9e-4U"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Minimize Energy\n",
        "print(\"Minimizing Energy...\")\n",
        "simulation.minimizeEnergy()\n",
        "\n",
        "simulation.reporters.append(PDBReporter('output.pdb', 1000))\n",
        "simulation.reporters.append(StateDataReporter(stdout, 1000, step=True,\n",
        "        potentialEnergy=True, temperature=True, volume=True))\n",
        "simulation.reporters.append(StateDataReporter(\"md_log.txt\", 100, step=True,\n",
        "        potentialEnergy=True, temperature=True, volume=True))"
      ],
      "metadata": {
        "id": "NlSgTQ-qe-1t"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "system.addForce(MonteCarloBarostat(1*bar, 300*kelvin))\n",
        "simulation.context.reinitialize(preserveState=True)"
      ],
      "metadata": {
        "id": "wEU_gtAkffEv"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Running NPT...\")\n",
        "simulation.step(10000)"
      ],
      "metadata": {
        "id": "QmuDunJ_flK1"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "print(\"Running NVT...\")\n",
        "simulation.step(10000)"
      ],
      "metadata": {
        "id": "toG_H7Msfndb"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "source": [
        "## Plotting"
      ],
      "metadata": {
        "id": "Z7hNfRtPf8nD"
      }
    },
    {
      "cell_type": "code",
      "source": [
        "## Plot data\n",
        "data = np.loadtxt(\"md_log.txt\", delimiter=',')\n",
        "\n",
        "step = data[:,0]\n",
        "potential_energy = data[:,1]\n",
        "temperature = data[:,2]\n",
        "volume = data[:,3]"
      ],
      "metadata": {
        "id": "TPw1FXkPf7-T"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Potential Energy\n",
        "plt.plot(step, potential_energy)\n",
        "plt.xlabel(\"Step\")\n",
        "plt.ylabel(\"Potential energy (kJ/mol)\")\n",
        "plt.savefig(\"potential_energy.png\")"
      ],
      "metadata": {
        "id": "UCoDnFKHf78O"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Temperature\n",
        "plt.plot(step, temperature)\n",
        "plt.xlabel(\"Step\")\n",
        "plt.ylabel(\"Temperature (K)\")\n",
        "plt.savefig(\"temperature.png\")"
      ],
      "metadata": {
        "id": "aGaj5q_Wf76B"
      },
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "source": [
        "## Volume\n",
        "plt.plot(step, volume)\n",
        "plt.xlabel(\"Step\")\n",
        "plt.ylabel(\"Volume (nm^3)\")\n",
        "plt.savefig(\"volume.png\")"
      ],
      "metadata": {
        "id": "WItY3vDof7f7"
      },
      "execution_count": null,
      "outputs": []
    }
  ]
}