MD_files_demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Running molecular dynamics simulations with `polymerist`\n",
    "In addition to useful polymer building and parameterization tools, polymerist also provides functionality which \\\n",
    "simplifies running molecular dynamics (MD) simulations by integrating with the OpenFF stack\n",
    "\n",
    "Here we will assume you have a parameterized structure for a single polymer availa, and will show how to :\n",
    "* Solvate the system (with Packmol backend)\n",
    "* Export the system to GROMACS, LAMMPS, and OpenMM (with OpenFF Interchange backend)\n",
    "* Reproducibly define and serialize parameters describing an MD simulation\n",
    "* Run a series of simulations defined by parameter sets (OpenMM only)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Logging and utilities setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Supressing annoying warnings (!must be done first!)\n",
    "import warnings\n",
    "warnings.catch_warnings(record=True)\n",
    "warnings.filterwarnings('ignore', category=UserWarning)\n",
    "warnings.filterwarnings('ignore', category=DeprecationWarning)\n",
    "\n",
    "# Logging\n",
    "import logging\n",
    "logging.basicConfig(level=logging.INFO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "def fetch_file(filename : str, extension : str, dir_to_search : Path=Path.cwd()) -> Path:\n",
    "    '''\n",
    "    Check a directory for the first file with a given name and extension\n",
    "    Return the file if found, or a message containing files with compatible \n",
    "    extension in the same directory otherwise\n",
    "    '''\n",
    "    available_files : dict[str, Path] = {\n",
    "        path.stem : path\n",
    "            for path in dir_to_search.glob(f'**/*.{extension}')\n",
    "    }\n",
    "    choice_str = ',\\n'.join(available_files.keys())\n",
    "\n",
    "    if filename in available_files:\n",
    "        return available_files[filename]\n",
    "    else:\n",
    "        raise ValueError(\n",
    "            f'No file called \"{filename}.{extension}\" found in {dir_to_search};' \\\n",
    "            f'\\nThe following (potentially) compatible files were found while searching:\\n{choice_str}'\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Loading a molecule\n",
    "We will employ pre-parameterized (i.e. with topologies, coordinates, and partial charges already assigned) polymer structures for this demo for compactness \\\n",
    "If interested in learning HOW to obtain such structures, see one of the other demos contained herein, such as [polymer_build_demo.ipynb](polymer_build_demo.ipynb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "STRUCT_DIR = Path('cleaned_structures') # this already should exist on import\n",
    "MD_DEMO_DIR = Path('MD_demo_files')\n",
    "MD_DEMO_DIR.mkdir(exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:rdkit:Enabling RDKit 2023.09.6 jupyter extensions\n",
      "INFO:numexpr.utils:Note: NumExpr detected 20 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 16.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 16 threads.\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7d3aa585a1c04130bcdc24bb166d667e",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "57d9e414f5684fdab03fad8b11c373a8",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "NGLWidget()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from polymerist.mdtools.openfftools.topology import topology_to_sdf, topology_from_sdf, get_largest_offmol\n",
    "from polymerist.genutils.fileutils.pathutils import assemble_path\n",
    "\n",
    "\n",
    "polymer_name = 'naturalrubber'\n",
    "\n",
    "polymer_sdf = fetch_file(polymer_name, extension='sdf', dir_to_search=STRUCT_DIR)\n",
    "polymer_outdir = MD_DEMO_DIR / polymer_name # only initialize this once an SDF has been found\n",
    "polymer_outdir.mkdir(exist_ok=True)\n",
    "\n",
    "polymer_topology = topology_from_sdf(polymer_sdf)\n",
    "polymer = get_largest_offmol(polymer_topology)\n",
    "polymer.visualize(backend='nglview')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define periodic box and pack that box with solvent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:polymerist.mdtools.openfftools.solvation.packing:Solvating 32.45169129382746 nm**3 Topology with 1082 water_TIP3P molecules to density of 0.997 g/(cm**3)\n",
      "INFO:polymerist.mdtools.openfftools.solvation.packing:Packmol packing converged\n",
      "INFO:polymerist.mdtools.openfftools.solvation.packing:Set solvated Topology box vectors to [[3.1878000319004056 0.0 0.0] [0.0 3.4547000169754027 0.0] [0.0 0.0 2.9467000007629394]] nanometer\n"
     ]
    }
   ],
   "source": [
    "from openmm.unit import gram, centimeter, nanometer\n",
    "\n",
    "from polymerist.mdtools.openfftools import boxvectors\n",
    "from polymerist.mdtools.openfftools.solvation import solvents \n",
    "from polymerist.mdtools.openfftools.solvation.packing import pack_topology_with_solvent\n",
    "\n",
    "\n",
    "box_padding = 1*nanometer # how far beyond the tight bounding box of the polymer to extend the periodic box\n",
    "solvent = solvents.water_TIP3P\n",
    "rho = 0.997 * gram / centimeter**3\n",
    "\n",
    "# calculate periodic box vectors\n",
    "box_dims    = boxvectors.get_topology_bbox(polymer_topology)\n",
    "box_vectors = boxvectors.box_vectors_flexible(box_dims)\n",
    "box_vectors = boxvectors.pad_box_vectors_uniform(box_vectors, box_padding)\n",
    "\n",
    "# calculate number of solvent molecules\n",
    "solvated_topology = pack_topology_with_solvent(polymer_topology, solvent=solvent, box_vecs=box_vectors, density=rho, exclusion=box_padding)\n",
    "solv_path = assemble_path(polymer_outdir, polymer_name, postfix=f'solv_{solvent.name}', extension='sdf')\n",
    "topology_to_sdf(solv_path, solvated_topology)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Selecting force field and generating Interchange object\n",
    "See [OpenFF Interchange documentation](https://docs.openforcefield.org/projects/interchange/en/stable/using/output.html) for details on how to output to various MD engine formats"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openff.toolkit import ForceField\n",
    "\n",
    "\n",
    "forcefield_names : list[str] = ['openff-2.2.0.offxml', 'tip3p.offxml']\n",
    "\n",
    "forcefield = ForceField(*forcefield_names)\n",
    "interchange = forcefield.create_interchange(solvated_topology, charge_from_molecules=[polymer])\n",
    "interchange.box = box_vectors"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running simulations using OpenMM\n",
    "`polymerist` and the OpenFF toolkit provide extensive support for running OpenMM simulations,\\\n",
    "due to the relative simplicity of automating these simulations thank to OpenMM's Pythonn API "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Defining reproducible and serializable simulation parameters\n",
    "These provide a means to exchange and cache information about how you set up a simulation \\\n",
    "These can be readily serialized to disc and are stored at the start of each simulation, if run through polymerist"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from openmm.unit import femtosecond, picosecond, nanosecond\n",
    "from openmm.unit import kelvin, atmosphere\n",
    "from polymerist.mdtools.openmmtools.parameters import (\n",
    "    SimulationParameters,\n",
    "    ThermoParameters,\n",
    "    IntegratorParameters,\n",
    "    ReporterParameters,\n",
    ")\n",
    "from polymerist.mdtools.openmmtools.reporters import DEFAULT_STATE_DATA_PROPS\n",
    "\n",
    "\n",
    "# define how data should be periodically reported during a simulation (will be shared between simulations in this examples)\n",
    "reporter_params = ReporterParameters( # these will be shared between both sets of parameters\n",
    "        report_trajectory=True,\n",
    "        traj_ext='dcd', # output to compressed binary trajectory files (recommended)\n",
    "        report_state_data=True,\n",
    "        state_data=DEFAULT_STATE_DATA_PROPS, # can tune these to taste\n",
    "        report_checkpoint=True, # also keep checkpoints of OpenMM objects (specific to Context and machine)\n",
    "        report_state=True,      # saving State is a bit redundant with checkpoints, but is machine-transferrable\n",
    "    )\n",
    "\n",
    "# define reproducible simulation parameter sets\n",
    "## 1) HIGH-TEMPERATURE ANNEAL FOR RELAXATION\n",
    "equil_params = SimulationParameters( \n",
    "    integ_params=IntegratorParameters(\n",
    "        time_step=1*femtosecond,\n",
    "        total_time=100*picosecond, # just a short simulation to demonstrate\n",
    "        num_samples=10, # don't want to take too many samples\n",
    "    ),\n",
    "    thermo_params=ThermoParameters(\n",
    "        ensemble='NPT', #\n",
    "        temperature=600*kelvin,\n",
    "        friction_coeff=1*picosecond**-1, # required for Langevin Thermostat\n",
    "        barostat_freq=25, # number of steps between barostat move attempts\n",
    "    ),\n",
    "    reporter_params=reporter_params,\n",
    ")\n",
    "\n",
    "## 2) LOW-TEMPERATURE PRODUCTION\n",
    "prod_params = SimulationParameters( \n",
    "    integ_params=IntegratorParameters(\n",
    "        time_step=2*femtosecond,  # NOTE that 2 fs step doesn't work well for unconstrainted FFs\n",
    "        total_time=0.5*nanosecond, # just a short simulation to demonstrate\n",
    "        num_samples=50, # don't want to take too many samples\n",
    "    ),\n",
    "    thermo_params=ThermoParameters(\n",
    "        ensemble='NVT', # clamp volume\n",
    "        temperature=300*kelvin,\n",
    "        friction_coeff=1*picosecond**-1, \n",
    "    ),\n",
    "    reporter_params=reporter_params,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Running a simulation \"schedule\"\n",
    "A \"schedule\" here denotes a serial sequence of simulations, each defined by their own SimulationParameters\n",
    "Note that subsequent parameter set in the schedule will generate a new directory containing the serailize parametesrs and all files output during the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "INFO:polymerist.mdtools.openmmtools.execution:Initializing simulation 1/2 (\"anneal\")\n",
      "INFO:polymerist.mdtools.openmmtools.thermo:Created MonteCarloBarostat Force(s) for NPT (Isothermal-isobaric) ensemble\n",
      "INFO:polymerist.mdtools.openmmtools.preparation:Added MonteCarloBarostat Force to System\n",
      "INFO:polymerist.mdtools.openmmtools.thermo:Created LangevinMiddleIntegrator for NPT (Isothermal-isobaric) ensemble\n",
      "WARNING:polymerist.mdtools.openmmtools.preparation:No valid State/State file provided, initializing State as None\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared DCDReporter which reports to MD_demo_files/naturalrubber/anneal/anneal_trajectory.dcd\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared CheckpointReporter which reports to MD_demo_files/naturalrubber/anneal/anneal_checkpoint.chk\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared StateReporter which reports to MD_demo_files/naturalrubber/anneal/anneal_state.xml\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared StateDataReporter which reports to MD_demo_files/naturalrubber/anneal/anneal_state_data.csv\n",
      "INFO:polymerist.mdtools.openmmtools.preparation:Setting positions in Context\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Performing energy minimization (initial PE = 56145.161448870145 kJ/mol)\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Energy successfully minimized (final PE = -53837.418308820284 kJ/mol)\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Saved energy-minimized Simulation Topology at MD_demo_files/naturalrubber/anneal/anneal_topology.pdb\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Saved serialized Simulation System at MD_demo_files/naturalrubber/anneal/anneal_system.xml\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Integrating 100 ps OpenMM Simulation for 100000 steps\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Simulation integration completed successfully\n",
      "INFO:polymerist.mdtools.openmmtools.execution:\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Initializing simulation 2/2 (\"production\")\n",
      "INFO:polymerist.mdtools.openmmtools.thermo:Created LangevinMiddleIntegrator for NVT (Canonical) ensemble\n",
      "WARNING:polymerist.mdtools.openmmtools.preparation:No valid State/State file provided, initializing State as None\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared DCDReporter which reports to MD_demo_files/naturalrubber/production/production_trajectory.dcd\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared CheckpointReporter which reports to MD_demo_files/naturalrubber/production/production_checkpoint.chk\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared StateReporter which reports to MD_demo_files/naturalrubber/production/production_state.xml\n",
      "INFO:polymerist.mdtools.openmmtools.reporters:Prepared StateDataReporter which reports to MD_demo_files/naturalrubber/production/production_state_data.csv\n",
      "INFO:polymerist.mdtools.openmmtools.preparation:Setting positions in Context\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Performing energy minimization (initial PE = 2.1096961101294497e+18 kJ/mol)\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Energy successfully minimized (final PE = -53702.92328318552 kJ/mol)\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Saved energy-minimized Simulation Topology at MD_demo_files/naturalrubber/production/production_topology.pdb\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Saved serialized Simulation System at MD_demo_files/naturalrubber/production/production_system.xml\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Integrating 0.5 ns OpenMM Simulation for 250000 steps\n",
      "INFO:polymerist.mdtools.openmmtools.execution:Simulation integration completed successfully\n",
      "INFO:polymerist.mdtools.openmmtools.execution:\n"
     ]
    }
   ],
   "source": [
    "from openff.interchange.interop.openmm._positions import to_openmm_positions\n",
    "\n",
    "from polymerist.mdtools.openmmtools.forcegroups import impose_unique_force_groups\n",
    "from polymerist.mdtools.openmmtools.execution import run_simulation_schedule\n",
    "\n",
    "\n",
    "# initialize core OpenMM objects\n",
    "omm_topology = interchange.to_openmm_topology()\n",
    "omm_system  = interchange.to_openmm(combine_nonbonded_forces=False)\n",
    "omm_positions = to_openmm_positions(interchange, include_virtual_sites=True)\n",
    "impose_unique_force_groups(omm_system) # ensure each Force is separate to enable mapping of energy contributions\n",
    "\n",
    "# run schedule\n",
    "schedule = { # simulations will be run in the order they appear here\n",
    "    'anneal'     : equil_params, \n",
    "    'production' : prod_params,\n",
    "}\n",
    "\n",
    "history = run_simulation_schedule(\n",
    "    polymer_outdir, \n",
    "    schedule, \n",
    "    init_top=omm_topology,\n",
    "    init_sys=omm_system,\n",
    "    init_pos=omm_positions,\n",
    "    return_history=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "polymerist-env",
   "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.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}