Skip to content

[WIP] OpenMMToolsMCEngine #1041

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
dwhswenson wants to merge 12 commits into
base: master
Choose a base branch
from
Draft

[WIP] OpenMMToolsMCEngine #1041

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
3a75a85
Add first draft of OpenMMToolsMCEngine
dwhswenson Jul 28, 2021
4290267
serialization, docstrings, mdtraj integration
dwhswenson Jul 31, 2021
d3b0537
start to tests on OpenMMToolsMCEngine
dwhswenson Jul 31, 2021
f4bcfd1
first few tests working
dwhswenson Jul 31, 2021
9942f71
more tests
dwhswenson Jul 31, 2021
311586e
that makes full test coverage!
dwhswenson Jul 31, 2021
a66b3e5
fix to tests when openmmtools not available
dwhswenson Jul 31, 2021
c2b51a5
copy openmmtools dicts when deserializing
dwhswenson Aug 2, 2021
4ee3d8b
add support for any kind of mover; more tests
dwhswenson Aug 3, 2021
207c0fb
Everything except last bits of example
dwhswenson Aug 4, 2021
116a9db
some cleanup to get the engine to run in example
dwhswenson Aug 7, 2021
e94fb60
update MC engine example
dwhswenson Aug 11, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Prev Previous commit
Next Next commit
Everything except last bits of example 
I seem to have a problem when I hit max length errors, though
  • Loading branch information
@dwhswenson
dwhswenson committed Aug 4, 2021
commit 207c0fb0a1d3db9fbd5df4ef175653a82efd4b0d
353 changes: 353 additions & 0 deletions examples/misc/mc_engine.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,353 @@
{
"cells": [
{
"cell_type": "markdown",
"id": "1ede3493",
"metadata": {},
"source": [
"# The OpenMMTools Monte Carlo Engine\n",
"\n",
"While it is very common to use molecular dynamics to generate the paths used in path sampling, a path can be any ordered sequence of snapshots. For example, one could generate a sequence of Markov Chain Monte Carlo steps connecting two states, and do path sampling on that.\n",
"\n",
"For those interested in doing this (within the domain of force-field based simulation), the OpenMMTools project has a subpackage for Markov Chain Monte Carlo, and OpenPathSampling has an engine that uses that subpackage.\n",
"\n",
"As usual, you set up the engine just as you would for a normal simulation with it. Then we wrap things from the underlying tool in the OPS engine wrapper."
]
},
{
"cell_type": "markdown",
"id": "a7f7e696",
"metadata": {},
"source": [
"## Creating an MCMC sampler with OpenMMTools"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "83c6a60e",
"metadata": {},
"outputs": [],
"source": [
"from simtk import unit\n",
"import openmmtools\n",
"from openmmtools import testsystems, cache, mcmc\n",
"from openmmtools.states import ThermodynamicState, SamplerState"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3fd3cd16",
"metadata": {},
"outputs": [],
"source": [
"testsystem = testsystems.AlanineDipeptideVacuum()\n",
"thermodynamic_state = ThermodynamicState(system=testsystem.system, \n",
" temperature=298*unit.kelvin)"
]
},
{
"cell_type": "markdown",
"id": "bd413a09",
"metadata": {},
"source": [
"In the OpenMMTools MCMC package, each move applies to all the atoms in its `atom_subset`. So to create a move that randomly selects a single atom and randomly displaces that atom, you need to create a displacement move for each atom, then join them in a `WeightedMove`."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e1ea14e2",
"metadata": {},
"outputs": [],
"source": [
"moves = [\n",
" mcmc.MCDisplacementMove(displacement_sigma=0.05*unit.angstrom,\n",
" atom_subset=[i])\n",
" for i in range(testsystem.mdtraj_topology.n_atoms)\n",
"]\n",
"move = mcmc.WeightedMove([(m, 1/len(moves)) for m in moves])"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "2baf56b5",
"metadata": {},
"outputs": [],
"source": [
"# use OpenMMTools to get good initial conditions\n",
"sampler_state = SamplerState(positions=testsystem.positions)\n",
"sampler = mcmc.MCMCSampler(thermodynamic_state, sampler_state, move)\n",
"sampler.minimize()"
]
},
{
"cell_type": "markdown",
"id": "ac95925b",
"metadata": {},
"source": [
"## Setting up path sampling with OpenPathSampling"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ffb97ab8",
"metadata": {},
"outputs": [],
"source": [
"import openpathsampling as paths\n",
"from openpathsampling.engines.openmm.mcengine import OpenMMToolsMCEngine, snapshot_from_sampler_state\n",
"\n",
"# we'll use the new storage, because it is faster\n",
"from openpathsampling.experimental.storage import monkey_patch_all, Storage\n",
"from openpathsampling.experimental.storage.collective_variables import MDTrajFunctionCV\n",
"paths = monkey_patch_all(paths)\n",
"\n",
"import mdtraj as md\n",
"import numpy as np"
]
},
{
"cell_type": "markdown",
"id": "403871a9",
"metadata": {},
"source": [
"### Creating the OPS engine\n",
"\n",
"The next two cells are the only ones specific to integrating this new engine type with OPS."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4ab50fe7",
"metadata": {},
"outputs": [],
"source": [
"mdtraj_topology = testsystem.mdtraj_topology\n",
"ops_topology = paths.engines.MDTrajTopology(mdtraj_topology)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "683a7c36",
"metadata": {},
"outputs": [],
"source": [
"engine = OpenMMToolsMCEngine(thermodynamic_state, move,\n",
" {'n_steps_per_frame': 10,\n",
" 'n_frames_max': 1000},\n",
" topology=mdtraj_topology)"
]
},
{
"cell_type": "markdown",
"id": "10bda885",
"metadata": {},
"source": [
"### Defining CVs and stable states"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "de025fc1",
"metadata": {},
"outputs": [],
"source": [
"# just to verify that we have the right atoms\n",
"[mdtraj_topology.atom(i) for i in [4, 6, 8, 14, 16]]"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "f6dc7342",
"metadata": {},
"outputs": [],
"source": [
"# CVs\n",
"phi = MDTrajFunctionCV(md.compute_dihedrals, topology=ops_topology, \n",
" period_min=-np.pi, period_max=np.pi,\n",
" indices=[[4, 6, 8, 14]]).named(\"phi\")\n",
"psi = MDTrajFunctionCV(md.compute_dihedrals, topology=ops_topology, \n",
" period_min=-np.pi, period_max=np.pi,\n",
" indices=[[6, 8, 14, 16]]).named(\"psi\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a8f499c8",
"metadata": {},
"outputs": [],
"source": [
"# TODO: check these values for Amber ff96\n",
"# estimates based on eyeballing https://doi.org/10.1098/rspa.2019.0036\n",
"# Mediocre state defs won't mess up the sampling, but aren't as helpful in teaching\n",
"C7eq = (\n",
" paths.PeriodicCVDefinedVolume(phi, lambda_min=-np.pi, lambda_max=-0.8,\n",
" period_min=-np.pi, period_max=np.pi)\n",
" & paths.PeriodicCVDefinedVolume(psi, lambda_min=0.5, lambda_max=3.5,\n",
" period_min=-np.pi, period_max=np.pi)\n",
").named(\"C7eq\")\n",
"C7ax = (\n",
" paths.PeriodicCVDefinedVolume(phi, lambda_min=0.5, lambda_max=1.5,\n",
" period_min=-np.pi, period_max=np.pi)\n",
" & paths.PeriodicCVDefinedVolume(psi, lambda_min=-2.0, lambda_max=-0.5,\n",
" period_min=-np.pi, period_max=np.pi)\n",
").named(\"C7ax\")\n",
"# period information in CVDefinedVolumes won't be required in OPS 2.0"
]
},
{
"cell_type": "markdown",
"id": "8a404464",
"metadata": {},
"source": [
"### Creating sampling network and move scheme\n",
"\n",
"For TPS, these are very easy. They get more complicated for TIS."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "6db89880",
"metadata": {},
"outputs": [],
"source": [
"network = paths.TPSNetwork(C7eq, C7ax).named(\"tps-network\")\n",
"scheme = paths.OneWayShootingMoveScheme(network, engine=engine).named('one-way TPS')"
]
},
{
"cell_type": "markdown",
"id": "8d856514",
"metadata": {},
"source": [
"### Getting an initial transition trajectory\n",
"\n",
"This is always one of the hardest parts of setting up a TPS simulation. In this example we will use a high temperature version of the engine to generate a transition from a long trajectory. This is not always the best way to generate an initial trajectory, but it works well enough for simple systems like this."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e239dab5",
"metadata": {},
"outputs": [],
"source": [
"# use the same move, different thermodynamic state, and allow more steps\n",
"hi_temp = engine = OpenMMToolsMCEngine(\n",
" thermodynamic_state=ThermodynamicState(system=testsystem.system, \n",
" temperature=700*unit.kelvin),\n",
" move=move,\n",
"# options={'n_steps_per_frame': 10, 'n_frames_max': 10000},\n",
" options={'n_steps_per_frame': 10, 'n_frames_max': 1000},\n",
" topology=mdtraj_topology\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ac7b19c3",
"metadata": {},
"outputs": [],
"source": [
"# make a snapshot from the minimized sampler state\n",
"snapshot = snapshot_from_sampler_state(sampler.sampler_state)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "4d92d442",
"metadata": {},
"outputs": [],
"source": [
"visit_all = paths.VisitAllStatesEnsemble([C7ax, C7eq])\n",
"traj = hi_temp.generate(snapshot, visit_all.can_append)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "1fe8dcb2",
"metadata": {},
"outputs": [],
"source": [
"# extract only the transition using scheme.initial_conditions_from_trajectories\n",
"init_conds = scheme.initial_conditions_from_trajectories(traj)"
]
},
{
"cell_type": "markdown",
"id": "7f0dca9a",
"metadata": {},
"source": [
"### Putting together the path sampling simulation"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c1be8dec",
"metadata": {},
"outputs": [],
"source": [
"storage = Storage(\"mc_tps.db\", mode='w')"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "ae09c265",
"metadata": {},
"outputs": [],
"source": [
"tps = paths.PathSamplingling(\n",
" storage=storage,\n",
" movescheme=scheme,\n",
" initial_conditions=init_conds\n",
")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "a04a0fb8",
"metadata": {},
"outputs": [],
"source": [
"tps.run(100)"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"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.7.10"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
9 changes: 9 additions & 0 deletions openpathsampling/engines/openmm/mcengine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -62,6 +62,15 @@ class OpenMMToolsMCEngine(paths.engines.DynamicsEngine):
Topology information for this sysetm. If an :class:`.MDTrajTopology`
is provided, this will enable MDTraj integration.
"""

base_snapshot_type = Snapshot

# taken from OpenMM engine; this should change
_default_options = {
'n_steps_per_frame': 10,
'n_frames_max': 5000,
}

def __init__(self, thermodynamic_state, move, options, topology=None):
# we assume that we have standard n_atoms * 3 dimensions
descriptor = paths.engines.SnapshotDescriptor.construct(
Expand Down
8 changes: 8 additions & 0 deletions openpathsampling/tests/test_openmm_mcengine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -136,6 +136,14 @@ def test_generate_next_frame_uninitialized_error(self, mcengine,
with pytest.raises(EngineNotInitializedError):
mcengine.generate_next_frame()

def test_generate_integration(self, mcengine, snapshot):
# check that integration with the generate method: we should be able
# to create a candidate trajectory for the ensemble
ensemble = paths.LengthEnsemble(5)
traj = mcengine.generate(snapshot, ensemble.can_append)
assert isinstance(traj, paths.Trajectory)
assert ensemble(traj)

def test_serialization_cycle(self, mcengine):
# if we serialize then deserialize, the to_dict of the deserialized
# object should be the same as the to_dict of the original
Expand Down