forked from SSAGESLabs/PySAGES
-
Notifications
You must be signed in to change notification settings - Fork 0
/
alanine-dipeptide_openmm.py
129 lines (93 loc) · 3.64 KB
/
alanine-dipeptide_openmm.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
#!/usr/bin/env python3
import matplotlib.pyplot as plt
import numpy
import pysages
# %%
from pysages.colvars import DihedralAngle
from pysages.methods import ABF
from pysages.utils import try_import
openmm = try_import("openmm", "simtk.openmm")
unit = try_import("openmm.unit", "simtk.unit")
app = try_import("openmm.app", "simtk.openmm.app")
# %%
pi = numpy.pi
adp_pdb = "../../inputs/alanine-dipeptide/adp-explicit.pdb"
T = 298.15 * unit.kelvin
dt = 2.0 * unit.femtoseconds
def generate_simulation(pdb_filename=adp_pdb, T=T, dt=dt):
pdb = app.PDBFile(pdb_filename)
ff = app.ForceField("amber99sb.xml", "tip3p.xml")
cutoff_distance = 1.0 * unit.nanometer
topology = pdb.topology
system = ff.createSystem(
topology, constraints=app.HBonds, nonbondedMethod=app.PME, nonbondedCutoff=cutoff_distance
)
# Set dispersion correction use.
forces = {}
for i in range(system.getNumForces()):
force = system.getForce(i)
forces[force.__class__.__name__] = force
forces["NonbondedForce"].setUseDispersionCorrection(True)
forces["NonbondedForce"].setEwaldErrorTolerance(1.0e-5)
positions = pdb.getPositions(asNumpy=True)
integrator = openmm.LangevinIntegrator(T, 1 / unit.picosecond, dt)
# platform = openmm.Platform.getPlatformByName(platform)
# simulation = app.Simulation(topology, system, integrator, platform)
simulation = app.Simulation(topology, system, integrator)
simulation.context.setPositions(positions)
simulation.minimizeEnergy()
return simulation
# %%
# helping function for plotting results
def plot_energy(result):
surface = numpy.asarray(result["free_energy"])
fig, ax = plt.subplots()
im = ax.imshow(
surface, interpolation="bicubic", origin="lower", extent=[-pi, pi, -pi, pi], aspect=1
)
ax.contour(surface, levels=15, linewidths=0.75, colors="k", extent=[-pi, pi, -pi, pi])
plt.colorbar(im)
fig.savefig("energy.pdf")
# %%
# %%
def plot_histogram(result):
surface = numpy.asarray(result["histogram"]) / numpy.nanmax(numpy.asarray(result["histogram"]))
fig, ax = plt.subplots()
im = ax.imshow(
surface, interpolation="bicubic", origin="lower", extent=[-pi, pi, -pi, pi], aspect=1
)
ax.contour(surface, levels=15, linewidths=0.75, colors="k", extent=[-pi, pi, -pi, pi])
plt.colorbar(im)
fig.savefig("histogram.pdf")
# %%
# %%
def plot_forces(result):
fig, ax = plt.subplots()
ax.set_xlabel("CV")
ax.set_ylabel("Forces $[\\epsilon]$")
forces = numpy.asarray(result["mean_force"])
x = numpy.asarray(result["mesh"])
plt.quiver(x, forces, width=(0.0002 * (x[x.shape[0] - 1, 0] - x[0, 0])), headwidth=3)
fig.savefig("forces.pdf")
# Stores forces and free energies for post-analysis
def save_energy_forces(result):
energy = numpy.asarray(result["free_energy"])
forces = numpy.asarray(result["mean_force"])
grid = numpy.asarray(result["mesh"])
numpy.savetxt("FES.csv", numpy.hstack([grid, energy.reshape(-1, 1)]))
numpy.savetxt("Forces.csv", numpy.hstack([grid, forces.reshape(-1, grid.shape[1])]))
def post_run_action(**kwargs):
kwargs.get("context").saveState("final.xml")
# %%
def main():
cvs = [DihedralAngle((4, 6, 8, 14)), DihedralAngle((6, 8, 14, 16))]
grid = pysages.Grid(lower=(-pi, -pi), upper=(pi, pi), shape=(32, 32), periodic=True)
method = ABF(cvs, grid, use_pinv=True)
raw_result = pysages.run(method, generate_simulation, 25, post_run_action=post_run_action)
result = pysages.analyze(raw_result, topology=(14,))
plot_energy(result)
plot_histogram(result)
save_energy_forces(result)
# %%
if __name__ == "__main__":
main()