Implementing Wales and Dolye's paper in OpenMM.
In [1]:
from math import ceil
import matplotlib.pyplot as pp
from itertools import islice
import numpy as np
from simtk.unit import angstroms
import simtk.openmm as mm
In [2]:
N_PARTICLES=23
system = mm.System()
force = mm.CustomNonbondedForce('4*(r6*r6 - r6); r6 = 1/(r*r*r*r*r*r);')
force.setNonbondedMethod(mm.CustomNonbondedForce.CutoffPeriodic)
system.setDefaultPeriodicBoxVectors((10,0,0), (0,10,0), (0,0,10))
force.setCutoffDistance(5)
for i in range(N_PARTICLES):
system.addParticle(1)
force.addParticle([])
system.addForce(force)
Out[2]:
In [3]:
integrator = mm.CustomIntegrator(0.001)
integrator.addGlobalVariable('a_cur', 0)
integrator.addGlobalVariable('a_old', 0)
integrator.addPerDofVariable('y_cur', 0)
integrator.addPerDofVariable('y_old', 0)
integrator.addComputeGlobal('a_cur', '0.5*(1+sqrt(1+(4*a_old*a_old)))')
integrator.addComputeGlobal('a_old', 'a_cur')
integrator.addComputePerDof('y_cur', 'x + dt*f')
integrator.addComputePerDof('y_old', 'y_cur')
integrator.addComputePerDof('x', 'y_cur + (a_old - 1) / a_cur * (y_cur - y_old)')
context = mm.Context(system, integrator)
In [4]:
def grid_coordinates(n_points):
"""Coordinates of grid points in 3d (e.g. (0,0,0), (0,0,1), ...)
Returns
-------
points : shape=(n_points, 3)
"""
griddim = ceil(n_points**(1.0/3.0))
mesh = np.indices((griddim, griddim, griddim))
points = np.array(list(zip(*[e.flatten() for e in mesh])))[:N_PARTICLES]
return points
def scatter3d(context):
xyz = context.getState(getPositions=True, enforcePeriodicBox=True).getPositions(asNumpy=True)._value
fig = pp.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(xyz[:,0], xyz[:,1], xyz[:,2])
return ax
In [5]:
def basin_hop(context, scaling=0.3, temperature=0.8, n_steps_per_hop=100, n_hops=1000):
state = context.getState(getPositions=True, getEnergy=True, enforcePeriodicBox=True)
positions = state.getPositions(asNumpy=True)._value
energy = state.getPotentialEnergy()._value
n_acceptances = 0
best_energy = 0
best_positions = None
for i in range(n_hops):
displacement = scaling * np.random.random(positions.shape)
context.setPositions(positions + displacement)
context.getIntegrator().step(n_steps_per_hop)
new_state = context.getState(getEnergy=True, getPositions=True, enforcePeriodicBox=True)
new_energy = new_state.getPotentialEnergy()._value
p = np.exp(-(new_energy - energy) / temperature)
if (np.random.rand() < p):
energy = new_energy
positions = new_state.getPositions(asNumpy=True)._value
n_acceptances += 1
if i % 1000 == 0 and i > 0:
print('step=%d, energy=%.3f, MC acceptance=%.3f' % (i, energy, float(n_acceptances) / i))
if energy < best_energy:
best_positions = positions
context.setPositions(best_positions)
mm.LocalEnergyMinimizer.minimize(context, 1e-6)
print('final energy', context.getState(getEnergy=True).getPotentialEnergy()._value)
print('Acceptance: %d/%d = %.3f' % (n_acceptances, n_hops, float(n_acceptances) / n_hops))
In [6]:
context.setPositions(grid_coordinates(N_PARTICLES))
basin_hop(context, n_hops=10000)
In [7]:
from simtk.openmm.app.element import carbon
from simtk.openmm.app import Topology, PDBFile
topology = Topology()
chain = topology.addChain()
residue = topology.addResidue('A', chain)
for i in range(N_PARTICLES):
topology.addAtom('XX', carbon, residue)
a,b,c = system.getDefaultPeriodicBoxVectors()
topology.setUnitCellDimensions((a[0]._value, b[1]._value, c[2]._value) * angstroms)
with open('out.pdb', 'w') as f:
PDBFile.writeFile(topology, context.getState(getPositions=True, enforcePeriodicBox=True).getPositions()._value, f)
In [9]:
print context.getState(getPositions=True, enforcePeriodicBox=True).getPositions(asNumpy=True)._value
In [ ]: