forked from MiaoLab20/gamd-openmm
-
Notifications
You must be signed in to change notification settings - Fork 0
/
Copy pathclean-gamd-runner.py
executable file
·178 lines (138 loc) · 7.16 KB
/
clean-gamd-runner.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
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
from simtk.openmm.app import *
from simtk.openmm import *
from simtk.unit import *
from sys import stdout
from sys import exit
import os
import sys
import time
import traceback
from gamd.langevin.total_boost_integrators import LowerBoundIntegrator
# from gamd.langevin.total_boost_integrators import UpperBoundIntegrator
# from gamd.langevin.dihedral_boost_integrators import import LowerBountIntegrator
# from gamd.langevin.dihedral_boost_integrators import import UpperBoundIntegrator
from gamd.stage_integrator import BoostType
from gamd.GamdLogger import GamdLogger
from gamd import utils as utils
import pprint
from gamd.langevin.total_boost_integrators import LowerBoundIntegrator as TotalBoostLowerBoundIntegrator
from gamd.langevin.total_boost_integrators import UpperBoundIntegrator as TotalBoostUpperBoundIntegrator
from gamd.langevin.dihedral_boost_integrators import LowerBoundIntegrator as DihedralBoostLowerBoundIntegrator
from gamd.langevin.dihedral_boost_integrators import UpperBoundIntegrator as DihedralBoostUpperBoundIntegrator
#
# NOTE: Don't do this. It moves the forces into separate groups, so that they don't get handled properly.
#
# for i, force in enumerate(system.getForces()):
# print(str(i) + " " + force.__class__.__name__)
# force.setForceGroup(i)
# if force.__class__.__name__ == 'PeriodicTorsionForce':
# group = i
def set_dihedral_group(system):
group = 1
for force in system.getForces():
if force.__class__.__name__ == 'PeriodicTorsionForce':
force.setForceGroup(group)
break
return group
def create_lower_total_boost_integrator(system):
return [set_dihedral_group(system), TotalBoostLowerBoundIntegrator()]
def create_upper_total_boost_integrator(system):
return [set_dihedral_group(system), TotalBoostUpperBoundIntegrator()]
def create_lower_dihedral_boost_integrator(system):
group = set_dihedral_group(system)
return [group, DihedralBoostLowerBoundIntegrator(group)]
def create_upper_dihedral_boost_integrator(system):
group = set_dihedral_group(system)
return [group, DihedralBoostUpperBoundIntegrator(group)]
def create_output_directories(directories):
for dir in directories:
os.makedirs(dir, 0o755)
def main():
starttime = time.time()
if len(sys.argv) == 1:
output_directory = "output"
else:
output_directory = sys.argv[1]
coordinates_file = './data/md-4ns.rst7'
prmtop_file = './data/dip.top'
restarting = False
restart_checkpoint_frequency = 1000
restart_checkpoint_filename = "gamd.backup"
if not restarting:
create_output_directories([output_directory, output_directory + "/states/", output_directory + "/positions/",
output_directory + "/pdb/", output_directory + "/checkpoints"])
# dihedral_boost = True
# (simulation, integrator) = createGamdSimulationFromAmberFiles(prmtop_file, coordinates_file, dihedral_boost=dihedral_boost)
prmtop = AmberPrmtopFile(prmtop_file)
inpcrd = AmberInpcrdFile(coordinates_file)
system = prmtop.createSystem(nonbondedMethod=PME, nonbondedCutoff=0.8*nanometer, constraints=HBonds)
# [group, integrator] = create_lower_total_boost_integrator(system)
[group, integrator] = create_upper_total_boost_integrator(system)
# [group, integrator] = create_lower_dihedral_boost_integrator(system)
# [group, integraotor] = create_upper_dihedral_boost_integrator(system)
number_of_steps_in_group = 1000
simulation = Simulation(prmtop.topology, system, integrator)
if restarting:
simulation.loadCheckpoint(restart_checkpoint_filename)
write_mode = "a"
start_step = int(integrator.getGlobalVariableByName("stepCount") // number_of_steps_in_group)
print("restarting from saved checkpoint:", restart_checkpoint_filename,
"at step:", start_step)
else:
simulation.context.setPositions(inpcrd.positions)
if inpcrd.boxVectors is not None:
simulation.context.setPeriodicBoxVectors(*inpcrd.boxVectors)
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(298.15*kelvin)
simulation.saveState(output_directory + "/states/initial-state.xml")
simulation.reporters.append(DCDReporter(output_directory + '/output.dcd', number_of_steps_in_group))
simulation.reporters.append(
utils.ExpandedStateDataReporter(system, output_directory + '/state-data.log', number_of_steps_in_group, step=True,
brokenOutForceEnergies=True, temperature=True, potentialEnergy=True,
totalEnergy=True, volume=True))
print("startup time:", time.time() - starttime)
write_mode = "w"
start_step = 1
gamd_logger = GamdLogger(output_directory + "/gamd.log", write_mode, integrator, simulation)
if not restarting:
gamd_logger.write_header()
print("Running " + str(integrator.get_total_simulation_steps()) + " steps")
for step in range(start_step, (integrator.get_total_simulation_steps() // number_of_steps_in_group) + 1):
if step % restart_checkpoint_frequency // number_of_steps_in_group == 0:
simulation.saveCheckpoint(restart_checkpoint_filename)
# TEST
# if step == 250 and not restarting:
# print("sudden, unexpected interruption!")
# exit()
# END TEST
gamd_logger.mark_energies(group)
try:
# print(integrator.get_current_state())
#
# NOTE: We need to save off the starting total and dihedral potential energies, since we
# end up modifying them by updating the state of the particles. This allows us to write
# out the values as they were at the beginning of the step for what all of the calculations
# for boosting were based on.
#
simulation.step(number_of_steps_in_group)
gamd_logger.write_to_gamd_log(step*number_of_steps_in_group)
# print(integrator.get_current_state())
except Exception as e:
print("Failure on step " + str(step * number_of_steps_in_group))
print(e)
# print(integrator.get_current_state())
gamd_logger.write_to_gamd_log(step)
sys.exit(2)
# simulation.loadCheckpoint('/tmp/testcheckpoint')
# debug_information = integrator.get_debugging_information()
# getGlobalVariableNames(integrator)
if step % integrator.ntave == 0:
# if step % 1 == 0:
simulation.saveState(output_directory + "/states/" + str(step * number_of_steps_in_group) + ".xml")
simulation.saveCheckpoint(output_directory + "/checkpoints/" + str(step * number_of_steps_in_group) + ".bin")
positions_filename = output_directory + '/positions/coordinates-' + str(step * number_of_steps_in_group) + '.csv'
integrator.create_positions_file(positions_filename)
# pp = pprint.PrettyPrinter(indent=2)
# pp.pprint(debug_information)
if __name__ == "__main__":
main()