demos/ANI_minimal.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "dbaceef4",
   "metadata": {},
   "source": [
    "# Small System, ~30 atoms\n",
    "\n",
    "ANI2x simulating a small peptide in vacuum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7aad2039",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import libraries\n",
    "from openmm.app import *\n",
    "from openmm import *\n",
    "from openmm.unit import *\n",
    "import openmmml\n",
    "import sys\n",
    "\n",
    "# Load small alanine tripeptide system\n",
    "pdb = PDBFile(\"../pdbs/aaa_capped.pdb\")\n",
    "potential = openmmml.MLPotential('ani2x')\n",
    "system = potential.createSystem(\n",
    "    pdb.topology,\n",
    "    nonbondedMethod=PME,\n",
    "    nonbondedCutoff=1*nanometer,\n",
    "    constraints=HBonds,\n",
    "    hydrogenMass=1*amu # required\n",
    ")\n",
    "\n",
    "# Setup simulation\n",
    "integrator = LangevinIntegrator(\n",
    "    300 * kelvin, \n",
    "    1 / picosecond, \n",
    "    2.0 * femtosecond,\n",
    ")\n",
    "simulation = Simulation(\n",
    "    pdb.topology,\n",
    "    system,\n",
    "    integrator,\n",
    ")\n",
    "STEPS = 1000\n",
    "simulation.reporters.append(\n",
    "    StateDataReporter(\n",
    "        sys.stdout, \n",
    "        30, \n",
    "        totalSteps=STEPS,\n",
    "        speed=True,\n",
    "        progress=True,\n",
    "        remainingTime=True,\n",
    "    )\n",
    ")\n",
    "simulation.context.setPositions(pdb.positions)\n",
    "\n",
    "# Minimise and run\n",
    "simulation.minimizeEnergy()\n",
    "simulation.step(STEPS)\n",
    "print(\"done\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04715fa6",
   "metadata": {},
   "source": [
    "# Large Solvated System, Mixed ANI/AMBER, ~2500 atoms\n",
    "\n",
    "ANI2x is used for the peptide, and AMBER for the solvent.\n",
    "\n",
    "ANI2x cannot be used alone for such a large system, as it uses too much memory. I read somewhere that memory usage is ~On^2 for number of atoms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21708b9e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Import libraries\n",
    "from openmm.app import *\n",
    "from openmm import *\n",
    "from openmm.unit import *\n",
    "import openmmml\n",
    "import sys\n",
    "\n",
    "# Load solvated, equilibrated system\n",
    "pdb = PDBFile(\"../pdbs_equilibrated/aaa/aaa_capped_amber_equilibrated.pdb\")\n",
    "\n",
    "# Select peptide (non-water) atoms to be simulated by ANI\n",
    "peptide_atoms = [atom.index for atom in pdb.topology.atoms() if atom.residue.name != \"HOH\"]\n",
    "\n",
    "# Create amber system\n",
    "amber_system = ForceField(\n",
    "    'amber14-all.xml',\n",
    "    'amber14/tip3p.xml'\n",
    ").createSystem(\n",
    "    pdb.topology,\n",
    "    nonbondedMethod=PME,\n",
    "    nonbondedCutoff=1*nanometer,\n",
    "    constraints=HBonds,\n",
    "    hydrogenMass=1*amu # required\n",
    ")\n",
    "\n",
    "# Setup mixed system\n",
    "potential = openmmml.MLPotential('ani2x')\n",
    "system = potential.createMixedSystem(\n",
    "    pdb.topology,\n",
    "    amber_system,\n",
    "    peptide_atoms\n",
    ")\n",
    "\n",
    "# Setup simulation\n",
    "integrator = LangevinIntegrator(\n",
    "    300 * kelvin, \n",
    "    1 / picosecond, \n",
    "    2 * femtosecond, # higher stepsizes are unstable\n",
    ")\n",
    "simulation = Simulation(\n",
    "    pdb.topology,\n",
    "    system,\n",
    "    integrator,\n",
    ")\n",
    "STEPS = 1000\n",
    "simulation.reporters.append(\n",
    "    StateDataReporter(\n",
    "        sys.stdout, \n",
    "        30, \n",
    "        totalSteps=STEPS,\n",
    "        speed=True,\n",
    "        progress=True,\n",
    "        remainingTime=True,\n",
    "    )\n",
    ")\n",
    "simulation.context.setPositions(pdb.positions)\n",
    "\n",
    "# Minimise and run\n",
    "simulation.minimizeEnergy()\n",
    "simulation.step(STEPS)\n",
    "print(\"done\")\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}