def add_reporters(self, simulation, outfn):
"Add reporters to a simulation"
def reporter_callback(report):
"""Callback for processing reporter output"""
self.log.info(report)
callback_reporter = CallbackReporter(reporter_callback,
self.report_interval,
step=True,
potentialEnergy=True,
temperature=True,
time=True,
total_steps=self.number_of_steps)
h5_reporter = HDF5Reporter(outfn,
self.report_interval,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True)
simulation.reporters.append(callback_reporter)
simulation.reporters.append(h5_reporter)
def evaluate(timestep, extra_chances):
xchmc = XCGHMCIntegrator(temperature=simulation_parameters['temperature'],
steps_per_hmc=1, timestep=timestep,
extra_chances=extra_chances, steps_per_extra_hmc=1,
collision_rate=1.0 / unit.picosecond)
from benchmark.testsystems import dhfr_constrained
testsystem = dhfr_constrained
platform = mm.Platform.getPlatformByName('CUDA')
platform.setPropertyDefaultValue('CudaPrecision', 'mixed')
platform.setPropertyDefaultValue('DeterministicForces', 'true')
testsystem.platform = platform
sim = testsystem.construct_simulation(xchmc)
# pick an equilibrated initial condition
testsystem.load_or_simulate_x_samples() # make sure x_samples is there...
sim.context.setPositions(testsystem.x_samples[-1])
from mdtraj.reporters import HDF5Reporter
name = "dhfr_xchmc_timestep={}fs, extra_chances={}".format(timestep.value_in_unit(unit.femtoseconds), extra_chances)
reporter = HDF5Reporter(file=name + ".h5", reportInterval=1000, velocities=True, cell=True)
sim.reporters.append(reporter)
from tqdm import tqdm
for _ in tqdm(range(10000)):
sim.context.setPositions(testsystem.sample_x_from_equilibrium())
sim.context.setVelocitiesToTemperature(simulation_parameters['temperature'])
sim.step(1)
print(xchmc.all_counts)
np.save(name + ".npy", xchmc.all_counts)
def run_equilibrium(system, topology, configuration, n_steps, report_interval, filename):
from mdtraj.reporters import HDF5Reporter
integrator = integrators.LangevinIntegrator()
simulation = app.Simulation(topology, system, integrator)
simulation.context.setPositions(configuration)
#equilibrate a little bit:
simulation.step(10000)
reporter = HDF5Reporter(filename, report_interval)
simulation.reporters.append(reporter)
simulation.step(n_steps)
def run_equilibrium(system, topology, configuration, n_steps, report_interval, equilibration_steps, filename):
from mdtraj.reporters import HDF5Reporter
integrator = integrators.LangevinIntegrator()
simulation = app.Simulation(topology, system, integrator)
simulation.context.setPositions(configuration)
openmm.LocalEnergyMinimizer.minimize(simulation.context)
#equilibrate:
integrator.step(equilibration_steps)
print("equilibration complete")
reporter = HDF5Reporter(filename, report_interval)
simulation.reporters.append(reporter)
simulation.step(n_steps)
def run_md_simulation(random_seed, simulation, pdb, args):
if args.SIM_RUN_SIMULATION:
print("Running simulation...")
if args.SIM_SET_INITIAL_VELOCITIES:
print(f" Setting up initial velocities at temperature {args.SIM_TEMP}")
simulation.context.setVelocitiesToTemperature(args.SIM_TEMP, random_seed)
reporting_to_screen_freq = max(1, int(round(args.SIM_N_STEPS / args.REP_STATE_N_SCREEN)))
reporting_to_file_freq = max(1, int(round(args.SIM_N_STEPS / args.REP_STATE_N_FILE)))
trajectory_freq = max(1, int(round(args.SIM_N_STEPS / args.TRJ_FRAMES)))
total_time = args.SIM_N_STEPS * args.SIM_TIME_STEP
print(" Number of steps: {} steps".format(args.SIM_N_STEPS))
print(" Time step: {}".format(args.SIM_TIME_STEP))
print(" Temperature: {}".format(args.SIM_TEMP))
print(" Total simulation time: {}".format(total_time.in_units_of(simtk.unit.nanoseconds)))
print(" Number of state reads: {} reads".format(args.REP_STATE_N_SCREEN))
print(" State reporting to screen every: {} step".format(reporting_to_screen_freq))
print(" State reporting to file every: {} step".format(reporting_to_file_freq))
print(" Number of trajectory frames: {} frames".format(args.TRJ_FRAMES))
print(" Trajectory frame every: {} step".format(trajectory_freq))
print(" Trajectory frame every: {}".format(trajectory_freq * args.SIM_TIME_STEP))
print(' Random seed:', random_seed)
print()
if args.TRJ_FILENAME_PDB:
simulation.reporters.append(PDBReporter(args.TRJ_FILENAME_PDB, trajectory_freq))
if args.TRJ_FILENAME_DCD:
simulation.reporters.append(DCDReporter(args.TRJ_FILENAME_DCD, trajectory_freq))
simulation.reporters.append(StateDataReporter(sys.stdout, reporting_to_screen_freq,
step=True, progress=True, potentialEnergy=True,
totalSteps=args.SIM_N_STEPS))
if args.REP_STATE_FILE_PATH:
simulation.reporters.append(StateDataReporter(args.REP_STATE_FILE_PATH, reporting_to_file_freq,
step=True, potentialEnergy=True))
if args.REP_STATE_FILE_H5_PATH:
simulation.reporters.append(HDF5Reporter(args.REP_STATE_FILE_H5_PATH, reporting_to_file_freq, velocities=True))
print('Running simulation...')
simulation.step(args.SIM_N_STEPS)
if args.TRJ_LAST_FRAME_PDB:
last_frame_file_name = args.TRJ_LAST_FRAME_PDB
state = simulation.context.getState(getPositions=True)
PDBFile.writeFile(pdb.topology, state.getPositions(), open(last_frame_file_name, 'w'))
if args.REP_PLOT_FILE_NAME:
plot_data(args.REP_STATE_FILE_PATH, args.REP_PLOT_FILE_NAME)
def equilibration_NPT(self, *args):
if not args:
simulation=self.simulation
positions=self.positions
# add barostat for NPT
self.system.addForce(omm.MonteCarloBarostat(simulation.pressure,
simulation.temperature*kelvin,
25))
simulation.context.setPositions(positions)
simulation.context.setVelocitiesToTemperature(simulation.temperature*kelvin)
# Define reporters
simulation.reporters.append(app.StateDataReporter(self.workdir,
simulation.trj_write['NPT'],
step=True,
potentialEnergy=True,
temperature=True,
progress=True,
remainingTime=True,
speed=True,
totalSteps=simulation.steps['NPT'],
separator='\t'))
simulation.reporters.append(HDF5Reporter(f'{self.workdir}/equilibration_NPT.h5',
simulation.trj_write['NPT'],
atomSubset=simulation.trj_indices))
print(simulation)
print('Restrained NPT equilibration...')
def run(opts):
system_xml_file = opts.system
integrator_xml_file = opts.integrator
coords_f = opts.coords
platform_name = opts.platform
deviceid = opts.device
write_freq = opts.write_freq
output = opts.output
nsteps = opts.nsteps
platform_properties = {
'OpenCLPrecision': 'mixed',
'OpenCLPlatformIndex': '0',
'OpenCLDeviceIndex': '0',
'CudaPrecision': 'mixed',
'CudaDeviceIndex': '0',
'CpuThreads': '1'
}
platform_properties['CudaDeviceIndex'] = deviceid
platform_properties['OpenCLDeviceIndex'] = deviceid
with open(system_xml_file, 'r') as f:
system = openmm.XmlSerializer.deserialize(f.read())
with open(integrator_xml_file, 'r') as f:
integrator = openmm.XmlSerializer.deserialize(f.read())
integrator.setRandomNumberSeed(random.randint(0, 2**16))
platform = openmm.Platform.getPlatformByName(platform_name)
properties = {
key: platform_properties[key]
for key in platform_properties
if key.lower().startswith(platform_name.lower())
}
if platform_name == 'CPU':
properties = {'CpuThreads': '1'}
print properties
# Create dummy topology to satisfy Simulation object
topology = app.Topology()
volume = wcadimer.natoms / wcadimer.density
length = volume**(1.0 / 3.0)
L = length.value_in_unit(units.nanometer)
topology.setUnitCellDimensions(Vec3(L, L, L) * units.nanometer)
simulation = app.Simulation(topology, system, integrator, platform,
properties)
init_data = np.load(coords_f)
coords = units.Quantity(init_data['coord'], units.nanometer)
simulation.context.setPositions(coords)
if 'veloc' in init_data:
velocs = units.Quantity(init_data['veloc'],
units.nanometer / units.picosecond)
simulation.context.setVelocities(velocs)
else:
simulation.context.setVelocitiesToTemperature(wcadimer.temperature)
# Attach reporters
simulation.reporters.append(
HDF5Reporter(output + '.h5', write_freq, atomSubset=[0, 1]))
simulation.reporters.append(
app.StateDataReporter(stdout,
20 * write_freq,
step=True,
potentialEnergy=True,
temperature=True,
progress=True,
remainingTime=True,
speed=True,
totalSteps=nsteps,
separator='\t'))
# Run segment
simulation.step(nsteps)
# Write restart data
state = simulation.context.getState(getPositions=True, getVelocities=True)
coords = state.getPositions(asNumpy=True)
velocs = state.getVelocities(asNumpy=True)
np.savez_compressed(output + '_restart.npz',
coords,
coord=coords,
veloc=velocs)
def test_reporter():
# stolen/modified from MDTraj/tests/test_reporter.py ... thanks rmcgibbo
tempdir = os.path.join(testdir, 'test_reporter')
os.makedirs(tempdir)
pdb = PDBFile(ref_file('ala2.pdb'))
forcefield = ForceField('amber99sbildn.xml', 'amber99_obc.xml')
system = forcefield.createSystem(pdb.topology,
nonbondedMethod=CutoffNonPeriodic,
nonbondedCutoff=1.0 * nanometers,
constraints=HBonds,
rigidWater=True)
integrator = LangevinIntegrator(300 * kelvin, 1.0 / picoseconds,
2.0 * femtoseconds)
integrator.setConstraintTolerance(0.00001)
platform = Platform.getPlatformByName('Reference')
simulation = Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.positions)
simulation.context.setVelocitiesToTemperature(300 * kelvin)
reffile = os.path.join(tempdir, 'traj.h5')
testfile = os.path.join(tempdir, 'traj-test.h5')
ref_reporter = HDF5Reporter(reffile,
2,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True,
velocities=True)
test_reporter = sample.MCReporter(testfile,
2,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True,
velocities=True)
simulation.reporters.append(ref_reporter)
simulation.reporters.append(test_reporter)
simulation.step(100)
ref_reporter.close()
test_reporter.close()
with HDF5TrajectoryFile(testfile) as f:
got = f.read()
yield lambda: eq(got.temperature.shape, (50, ))
yield lambda: eq(got.potentialEnergy.shape, (50, ))
yield lambda: eq(got.kineticEnergy.shape, (50, ))
yield lambda: eq(got.coordinates.shape, (50, 22, 3))
yield lambda: eq(got.velocities.shape, (50, 22, 3))
yield lambda: eq(got.time, 0.002 * 2 * (1 + np.arange(50)))
yield lambda: f.topology == mdtraj.load(ref_file('ala2.pdb')).top
ref_traj = mdtraj.load(reffile)
test_traj = mdtraj.load(testfile)
yield lambda: eq(ref_traj.xyz, test_traj.xyz)
yield lambda: eq(ref_traj.unitcell_vectors, test_traj.unitcell_vectors)
yield lambda: eq(ref_traj.time, test_traj.time)
yield lambda: eq(ref_traj.xyz, test_traj.xyz)
yield lambda: eq(ref_traj.unitcell_vectors, test_traj.unitcell_vectors)
def SMD(system, prmtop, platform, platformProperties, temperature, positions, velocities, keyInteraction, spring_k, dist_in, dist_fin, SMD_num, save_step, move_force_step):
# See page 456 of http://ambermd.org/doc12/Amber17.pdf
pullforce = mm.CustomExternalForce('k_sp*0.5*(dx^2+dy^2+dz^2); \
dx=x-(x0+displace_x); \
dy=y-(y0+displace_y); \
dz=z-(z0+displace_z);')
pullforce.addPerParticleParameter('k_sp')
pullforce.addPerParticleParameter('x0')
pullforce.addPerParticleParameter('y0')
pullforce.addPerParticleParameter('z0')
pullforce.addGlobalParameter("displace_x", 0.0 * u.nanometer)
pullforce.addGlobalParameter("displace_y", 0.0 * u.nanometer)
pullforce.addGlobalParameter("displace_z", 0.0 * u.nanometer)
keyInteraction_pos = [positions[keyInteraction[0]], positions[keyInteraction[1]]]
keyInteraction_dist = np.linalg.norm(keyInteraction_pos[0] - keyInteraction_pos[1])
keyInteraction_vect = (keyInteraction_pos[1] - keyInteraction_pos[0]) / keyInteraction_dist
keyInteraction_vect = u.Quantity(value=keyInteraction_vect, unit=u.nanometers)
pullto = keyInteraction_pos[0] + 0.25 * keyInteraction_vect
pullto_old = pullto
pullforce.addParticle(keyInteraction[1], [spring_k, pullto[0], pullto[1], pullto[2] ])
system.addForce(pullforce)
integrator = mm.LangevinIntegrator(temperature, 4/u.picosecond, 0.002*u.picosecond)
simulation = app.Simulation(prmtop.topology, system, integrator, platform,
platformProperties)
simulation.context.setPositions(positions)
simulation.context.setVelocities(velocities)
force_val_old = -spring_k*(keyInteraction_dist - dist_in)
energy_val_old = u.Quantity(value=0, unit=u.kilocalorie/u.mole)
f=open('duck_'+str(temperature).split()[0].replace('.0','K')+'_'+str(SMD_num)+'.dat','w')
steps = int((dist_fin.value_in_unit(u.nanometer) / 0.000001 - dist_in.value_in_unit(u.nanometer) / 0.000001)) / move_force_step
pull_distance = 0.000001 * move_force_step
#write trajectory
top = md.load_prmtop('system_solv.prmtop')
atom_subset = top.select('not water')
simulation.reporters.append(app.StateDataReporter("smd_"+str(temperature).split()[0].replace('.0','K')+"_"+str(SMD_num)+".csv", move_force_step, step=True, time=True, totalEnergy=True, kineticEnergy=True, potentialEnergy=True,
temperature=True, density=True, progress=True, totalSteps=move_force_step*steps, speed=True))
simulation.reporters.append(HDF5Reporter("smd_"+str(temperature).split()[0].replace('.0','K')+"_"+str(SMD_num)+".h5", move_force_step*20, atomSubset=atom_subset))
for i in range(steps):
state = simulation.context.getState(getPositions=True)
pos_keyInt = state.getPositions()
keyInteraction_pos = [pos_keyInt[keyInteraction[0]], pos_keyInt[keyInteraction[1]]]
keyInteraction_dist = np.linalg.norm(keyInteraction_pos[0]-keyInteraction_pos[1])
keyInteraction_vect = (keyInteraction_pos[1] - keyInteraction_pos[0]) / keyInteraction_dist
keyInteraction_vect = u.Quantity(value=keyInteraction_vect, unit=u.nanometers)
pullto = keyInteraction_pos[0] + (0.25 + float(i) * pull_distance) * keyInteraction_vect
displace = pullto - pullto_old
simulation.context.setParameter('displace_x', displace[0])
simulation.context.setParameter('displace_y', displace[1])
simulation.context.setParameter('displace_z', displace[2])
if i == 0:
distance = 0.0
else:
distance = pull_distance
dist_spring = (0.25 + float(i) * pull_distance) * u.nanometer
force_val = -spring_k * (keyInteraction_dist - dist_spring)
energy_val = energy_val_old + (distance * u.nanometer) * 0.5 * (force_val+force_val_old)
force_val_old = force_val
energy_val_old = energy_val
if (i%int(save_step/move_force_step)) == 0:
f.write(str(i)+' '+str(dist_spring)+' '+str(keyInteraction_dist)+' '+str(force_val)+' '+str(energy_val)+'\n')
simulation.step(move_force_step)
#f.write(str(i)+' '+str(dist_spring)+' '+str(keyInteraction_dist)+' '+str(force_val)+' '+str(energy_val)+'\n')
f.close()
def via_openmm(cls,
parmed_obj,
file_name,
file_path="./",
platform="CUDA",
num_steps=5000 * 500,
write_out_freq=5000,
report_equilibration=False,
report_production=False,
constrain_all_bonds=True,
**kwargs):
"""
Runs simulation using OpenMM.
Parameters
------------
parmed_obj : parmed.structure
Parmed object of the fully parameterised simulated system.
file_name : str
No file type postfix is necessary
file_path : str
Default to current directory
platform : str
The computing architecture to do the calculation, default to CUDA, CPU, OpenCL is also possible.
num_steps : int
Number of production simulation to run, default 2,500,000 steps, i.e. 5 ns.
write_out_freq : int
Write out every nth frame of simulated trajectory, default to every 5000 frame write out one, i.e. 10 ps per frame.
Returns
--------
path : str
The absolute path where the trajectory is written to.
"""
platform = Platform.getPlatformByName(platform)
pmd = parmed_obj
path = '{}/{}.h5'.format(file_path, file_name)
constrain_what_bond = app.AllBonds if constrain_all_bonds else app.HBonds
##################3
system = pmd.createSystem(nonbondedMethod=app.PME,
nonbondedCutoff=1 * unit.nanometer,
constraints=constrain_what_bond)
# tmp_dir = tempfile.mkdtemp(dir = file_path) #FIXME just change this to debug dir in dubug mode
# pmd.save("{}/tmp.inpcrd".format(tmp_dir), overwrite = True)
# inpcrd = app.AmberInpcrdFile("{}/tmp.inpcrd".format(tmp_dir))
# pmd.save("{}/tmp.prmtop".format(tmp_dir), overwrite = True)
# prmtop = app.AmberPrmtopFile("{}/tmp.prmtop".format(tmp_dir))
# system = prmtop.createSystem(nonbondedMethod=app.PME, nonbondedCutoff=1*unit.nanometer, constraints=constrain_what_bond)
#####################3
thermostat = AndersenThermostat(cls.temperature, 1 / unit.picosecond)
system.addForce(thermostat)
barostat = MonteCarloBarostat(cls.pressure, cls.temperature)
system.addForce(barostat)
integrator = VerletIntegrator(cls.time_step)
# integrator = LangevinMiddleIntegrator(cls.temperature, 1/unit.picosecond, cls.time_step)
# simulation = Simulation(prmtop.topology, system, integrator)
# simulation.context.setPositions(inpcrd.positions)
# simulation.context.setPeriodicBoxVectors(*inpcrd.boxVectors)
simulation = Simulation(pmd.topology, system, integrator, platform)
simulation.context.setPeriodicBoxVectors(*pmd.box_vectors)
simulation.context.setPositions(pmd.positions)
simulation.minimizeEnergy()
#Eq
try:
cls.equil_steps = kwargs["equil_steps"]
except KeyError:
pass
if report_equilibration:
#print(cls.equil_steps, " steps")
simulation.reporters.append(
StateDataReporter("{}/equilibration_{}.dat".format(
file_path, file_name),
cls.equil_steps // 5000,
step=True,
volume=True,
temperature=True))
simulation.step(cls.equil_steps)
state = simulation.context.getState(getPositions=True,
getVelocities=True)
pmd.positions, pmd.velocities, pmd.box_vectors = state.getPositions(
), state.getVelocities(), state.getPeriodicBoxVectors()
#Production
del system
del simulation
system = pmd.createSystem(nonbondedMethod=app.PME,
nonbondedCutoff=1 * unit.nanometer,
constraints=constrain_what_bond)
thermostat = AndersenThermostat(cls.temperature, 1 / unit.picosecond)
system.addForce(thermostat)
#barostat = MonteCarloBarostat(1.013 * unit.bar, 298.15 * unit.kelvin)
#system.addForce(barostat)
integrator = VerletIntegrator(cls.time_step)
simulation = Simulation(pmd.topology, system, integrator, platform)
simulation.context.setPeriodicBoxVectors(*pmd.box_vectors)
simulation.context.setPositions(pmd.positions)
if report_production:
simulation.reporters.append(
StateDataReporter("{}/production_{}.dat".format(
file_path, file_name),
num_steps // 50000,
step=True,
potentialEnergy=True,
temperature=True))
simulation.reporters.append(HDF5Reporter(path, write_out_freq))
simulation.step(num_steps)
return os.path.abspath(path)
def test_reporter_subset():
tempdir = os.path.join(dir, 'test2')
os.makedirs(tempdir)
pdb = PDBFile(get_fn('native2.pdb'))
pdb.topology.setUnitCellDimensions([2, 2, 2])
forcefield = ForceField('amber99sbildn.xml', 'amber99_obc.xml')
system = forcefield.createSystem(pdb.topology,
nonbondedMethod=CutoffPeriodic,
nonbondedCutoff=1 * nanometers,
constraints=HBonds,
rigidWater=True)
integrator = LangevinIntegrator(300 * kelvin, 1.0 / picoseconds,
2.0 * femtoseconds)
integrator.setConstraintTolerance(0.00001)
platform = Platform.getPlatformByName('Reference')
simulation = Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.positions)
simulation.context.setVelocitiesToTemperature(300 * kelvin)
hdf5file = os.path.join(tempdir, 'traj.h5')
ncfile = os.path.join(tempdir, 'traj.nc')
dcdfile = os.path.join(tempdir, 'traj.dcd')
atomSubset = [0, 1, 2, 4, 5]
reporter = HDF5Reporter(hdf5file,
2,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True,
velocities=True,
atomSubset=atomSubset)
reporter2 = NetCDFReporter(ncfile,
2,
coordinates=True,
time=True,
cell=True,
atomSubset=atomSubset)
reporter3 = DCDReporter(dcdfile, 2, atomSubset=atomSubset)
simulation.reporters.append(reporter)
simulation.reporters.append(reporter2)
simulation.reporters.append(reporter3)
simulation.step(100)
reporter.close()
reporter2.close()
reporter3.close()
t = md.load(get_fn('native.pdb'))
t.restrict_atoms(atomSubset)
with HDF5TrajectoryFile(hdf5file) as f:
got = f.read()
eq(got.temperature.shape, (50, ))
eq(got.potentialEnergy.shape, (50, ))
eq(got.kineticEnergy.shape, (50, ))
eq(got.coordinates.shape, (50, len(atomSubset), 3))
eq(got.velocities.shape, (50, len(atomSubset), 3))
eq(got.cell_lengths, 2 * np.ones((50, 3)))
eq(got.cell_angles, 90 * np.ones((50, 3)))
eq(got.time, 0.002 * 2 * (1 + np.arange(50)))
assert f.topology == md.load(get_fn('native.pdb'),
atom_indices=atomSubset).topology
with NetCDFTrajectoryFile(ncfile) as f:
xyz, time, cell_lengths, cell_angles = f.read()
eq(cell_lengths, 20 * np.ones((50, 3)))
eq(cell_angles, 90 * np.ones((50, 3)))
eq(time, 0.002 * 2 * (1 + np.arange(50)))
eq(xyz.shape, (50, len(atomSubset), 3))
hdf5_traj = md.load(hdf5file)
dcd_traj = md.load(dcdfile, top=hdf5_traj)
netcdf_traj = md.load(ncfile, top=hdf5_traj)
# we don't have to convert units here, because md.load already handles
# that
eq(hdf5_traj.xyz, netcdf_traj.xyz)
eq(hdf5_traj.unitcell_vectors, netcdf_traj.unitcell_vectors)
eq(hdf5_traj.time, netcdf_traj.time)
eq(dcd_traj.xyz, hdf5_traj.xyz)
eq(dcd_traj.unitcell_vectors, hdf5_traj.unitcell_vectors)
def production_NPT(self, protocol):
prod_trj=f'{self.workdir}/production_NPT'
self.simulation.context.setPositions(self.positions)
# =============================================================================
# NO need for now since NPT eq protocols are not removing the barostat
# # add MC barostat for NPT
# self.system.addForce(omm.MonteCarloBarostat(self.pressure,
# self.temperature,
# 25))
# #TODO: force is hardcoded, make it go away. Check units throughout!
# =============================================================================
# Define reporters
#TODO: Decide on wether same extent as steps for reporter
#TODO: Link to mlflow or streamz
#TODO: Use append on DCD for control of continuation
self.simulation.reporters.append(DCDReporter(f'{prod_trj}.dcd', protocol['report']))
self.simulation.reporters.append(HDF5Reporter(f'{prod_trj}.h5', protocol['report'], atomSubset=self.trj_indices))
self.simulation.reporters.append(app.StateDataReporter(f'{prod_trj}.csv',
protocol['report'],
step=True,
totalEnergy=True,
temperature=True,
density=True,
progress=True,
remainingTime=True,
speed=True,
totalSteps=protocol['step'],
separator='\t'))
positions_first = self.simulation.context.getState(getPositions=True, getVelocities=True).getPositions()
self.writePDB(self.simulation.topology, positions_first, name='production_NPT_0')
#############
## WARNIN ##
#############
trj_time=protocol['step']*self.dt
print(f"NPT production ({trj_time})...")
self.simulation.step(protocol['step'])
#############
## WARNOUT ##
#############
state=self.simulation.context.getState(getPositions=True, getVelocities=True, getEnergy=True)
self.positions = state.getPositions()
self.velocities = state.getVelocities()
self.energy=state.getPotentialEnergy()
print('NPT production finished.')
self.EQ_NPT_PDB=self.writePDB(self.simulation.topology, self.positions, name='production_NPT')
Eo=self.simulation.context.getState(getEnergy=True).getPotentialEnergy()
print(f'System energy: {Eo}')
return self.simulation
def simulationSpawner(self, protocol, label='NPT', kind='equilibration', index=0):
self.energy=self.simulation.context.getState(getEnergy=True).getTotalEnergy()
print(f'Spwaning new simulation:')
print(f'\tSystem total energy: {self.energy}')
if kind == 'minimization':
self.simulation.minimizeEnergy()
self.structures['minimization']=self.writePDB(self.topology, self.positions, name='minimization')
print(f'\tPerforming energy minimization: {Ei}')
else:
name=f'{kind}_{label}'
trj_file=f'{self.workdir}/{kind}_{label}'
#If there are user defined costum forces
try:
if protocol['restrained_sets']:
self.system=self.setRestraints(protocol['restrained_sets'])
except KeyError:
pass
if label == 'NPT':
#TODO: frequency is hardcoded to 25, make it go away. Check units throughout!
print(f'\tAdding MC barostat for {name} ensemble')
self.system.addForce(omm.MonteCarloBarostat(self.pressure, self.temperature, 25))
self.simulation.context.setPositions(self.positions)
if index == 1 and kind == 'equilibration':
print(f'\tFirst equilibration run {name}. Assigning velocities to {self.temperature}')
self.simulation.context.setVelocitiesToTemperature(self.temperature)
# Define reporters
self.simulation.reporters.append(DCDReporter(f'{trj_file}.dcd', protocol['report']))
self.simulation.reporters.append(HDF5Reporter(f'{trj_file}.h5', protocol['report'], atomSubset=self.trj_indices))
self.simulation.reporters.append(app.StateDataReporter(f'{trj_file}.csv',
protocol['report'],
step=True,
totalEnergy=True,
temperature=True,
density=True,
progress=True,
remainingTime=True,
speed=True,
totalSteps=protocol['step'],
separator='\t'))
positions_first = self.simulation.context.getState(getPositions=True).getPositions()
self.structures['f{name}_ref']=self.writePDB(self.simulation.topology, positions_first, name=f'{name}_0')
#############
## WARNIN ##
#############
trj_time=protocol['step']*self.dt
print(f'\t{kind} simulation in {label} ensemble ({trj_time})...')
self.simulation.step(protocol['step'])
#############
## WARNOUT ##
#############
print(f'\t{kind} simulation in {label} ensemble ({trj_time}) completed.')
self.structures[f'{name}']=self.writePDB(self.simulation.topology, self.positions, name='{name}')
state=self.simulation.context.getState(getPositions=True, getVelocities=True, getEnergy=True)
self.positions = state.getPositions()
self.velocities = state.getVelocities()
self.energy=state.getPotentialEnergy()
#Remove Costum forces (always initialized in next run)
if label == 'NPT':
self.system=self.forceHandler(self.system, kinds=['CustomExternalForce', 'MonteCarloBarostat'])
elif label == 'NVT':
self.system=self.forceHandler(self.system, kinds=['CustomExternalForce'])
return self.simulation
molsys = msm.build.add_terminal_cappings(molsys, N_terminal='ACE', C_terminal='NME')
molsys = msm.build.add_hydrogens(molsys, pH=7.4)
_ = msm.convert(molsys, to_form='vacuum.msmpk')
# solvated
print('Solvated system in msmpk file...')
molsys = msm.build.solvate([molsys, {'forcefield':'AMBER14', 'water_model':'TIP3P'}],
box_geometry='truncated octahedral', clearance='14.0 angstroms',
to_form='molsysmt.MolSys', engine="OpenMM", verbose=False)
_ = msm.convert(molsys, to_form='solvated.msmpk')
# simulation
print('Trajectory files...')
modeller = msm.convert(molsys, to_form='openmm.Modeller')
forcefield = app.ForceField("amber14-all.xml", "amber14/tip3p.xml")
system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.PME, nonbondedCutoff=1.2*unit.nanometer, constraints=app.HBonds)
integrator = mm.LangevinIntegrator(300*unit.kelvin, 1.0/unit.picosecond, 2.0*unit.femtoseconds)
platform = mm.Platform.getPlatformByName('CUDA')
simulation = app.Simulation(modeller.topology, system, integrator, platform)
simulation.context.setPositions(modeller.positions)
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(300*unit.kelvin)
simulation.reporters.append(app.StateDataReporter(stdout, 50000, progress=True,
potentialEnergy=True, temperature=True, remainingTime=True, totalSteps=1000000))
simulation.reporters.append(app.DCDReporter('traj_explicit_solvent.dcd', 50000, enforcePeriodicBox=True))
simulation.reporters.append(HDF5Reporter('traj_explicit_solvent.h5', 50000))
simulation.step(1000000)
simulation.reporters[2].close()
final_positions = simulation.context.getState(getPositions=True).getPositions()
modeller = app.Modeller(pdb.topology, pdb.positions)
modeller.addSolvent(forcefield, model='tip3p', boxSize=BOX_SIZE, ionicStrength=40*unit.millimolar)
system = forcefield.createSystem(modeller.topology, nonbondedMethod=app.PME,
nonbondedCutoff=9.5*unit.angstroms,
constraints=app.HBonds, rigidWater=True,
ewaldErrorTolerance=0.0005)
system.addForce(mm.MonteCarloBarostat(1*unit.atmosphere, TEMPERATURE, 25))
integrator = mm.LangevinIntegrator(TEMPERATURE, 1.0/unit.picoseconds, TIMESTEP)
integrator.setConstraintTolerance(0.00001)
platform = mm.Platform.getPlatformByName('CUDA')
properties = {'CudaPrecision': 'mixed'}
simulation = app.Simulation(modeller.topology, system, integrator, platform, properties)
simulation.context.setPositions(modeller.positions)
print 'Setting Velocities'
simulation.context.setVelocitiesToTemperature(TEMPERATURE)
print 'Adding reporters'
simulation.reporters.append(HDF5Reporter(OUT_FN, REPORT_INTERVAL))
simulation.reporters.append(ProgressReporter(sys.stdout, REPORT_INTERVAL, N_STEPS))
state = simulation.context.getState(getPositions=True, getEnergy=True)
for r in simulation.reporters:
r.report(simulation, state)
simulation.step(N_STEPS)
md_integrator = openmm.openmm.LangevinIntegrator(temperature, 1/unit.picosecond, 0.002*unit.picoseconds)
print('minizmizing')
#select if min is not presnet, otherwise 0
if 1:
md_simulation = openmm.app.simulation.Simulation(topology=pdb.topology, system=system, integrator=md_integrator)
md_simulation.context.setPositions(pdb.positions)
openmm.LocalEnergyMinimizer.minimize(md_simulation.context)
md_simulation.saveState('min.xml')
else:
md_simulation = openmm.app.simulation.Simulation(topology=pdb.topology, system=system, integrator=md_integrator, state='min.xml')
md_simulation.context.setVelocitiesToTemperature(temperature)
md_simulation.reporters.append(openmm.app.dcdreporter.DCDReporter('npt.dcd', 100000))
md_simulation.reporters.append(HDF5Reporter('npt.h5', 100000))
md_simulation.reporters.append(openmm.app.statedatareporter.StateDataReporter('nptinfo.csv', 100000, step=True, potentialEnergy=True, kineticEnergy=True, totalEnergy=True, temperature=True, volume=True))
if 1:
print('Equilibrating Volume')
systemFromContext = md_simulation.context.getSystem()
print systemFromContext.getDefaultPeriodicBoxVectors()
md_simulation.step(stepsNPT)
temp_xml = 'temp.xml'
barostat_index = systemFromContext.getNumForces() -1
md_simulation.system.removeForce(barostat_index)
# md_simulation.saveState('temp.xml')
stateinfo = md_simulation.context.getState(True, True, True, True, True, True)
statepos = stateinfo.getPositions(asNumpy=True)[:]
1 * bar, 0.0 * bar * nanometer, 308 * kelvin,
MonteCarloMembraneBarostat.XYIsotropic, MonteCarloMembraneBarostat.ZFree,
15)
system_generator = SystemGenerator(
forcefields=['amber/lipid17.xml', 'amber/tip3p_standard.xml'],
small_molecule_forcefield='gaff-2.11',
barostat=membrane_barostat,
forcefield_kwargs=forcefield_kwargs,
periodic_forcefield_kwargs=periodic_forcefield_kwargs)
system = system_generator.create_system(pdb.topology, molecules=molecule)
integrator = LangevinIntegrator(300 * kelvin, 1 / picosecond,
0.002 * picosecond)
platform = Platform.getPlatformByName('CUDA')
simulation = app.Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.positions)
simulation.loadState('parent.xml')
simulation.reporters.append(
StateDataReporter('seg.nfo',
5000,
step=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True))
simulation.reporters.append(HDF5Reporter('seg.h5', 10000))
simulation.step(50000)
simulation.saveState('seg.xml')
def equilibration_NVT(self, protocol, index=0):
eq_trj=f'{self.workdir}/equilibration_NVT'
if protocol['restrained_sets']:
#call to setRestraints, returns updated system. Check protocol['restrained_sets'] definitions on setSimulation.
self.system=self.setRestraints(protocol['restrained_sets'])
# =============================================================================
# self.system.addForce(omm.MonteCarloBarostat(self.pressure,
# self.temperature,
# 25))
# #TODO: force is hardcoded, make it go away. Check units throughout!
# =============================================================================
self.simulation.context.setPositions(self.positions)
if index == 0:
self.simulation.context.setVelocitiesToTemperature(self.temperature)
# Define reporters
self.simulation.reporters.append(DCDReporter(f'{eq_trj}.dcd', protocol['report']))
self.simulation.reporters.append(HDF5Reporter(f'{eq_trj}.h5', protocol['report'], atomSubset=self.trj_indices))
self.simulation.reporters.append(app.StateDataReporter(f'{eq_trj}.csv',
protocol['report'],
step=True,
totalEnergy=True,
temperature=True,
density=True,
progress=True,
remainingTime=True,
speed=True,
totalSteps=protocol['step'],
separator='\t'))
positions_first = self.simulation.context.getState(getPositions=True).getPositions()
self.writePDB(self.simulation.topology, positions_first, name='equilibration_NVT_0')
#############
## WARNIN ##
#############
trj_time=protocol['step']*self.dt
print(f"Restrained NVT equilibration ({trj_time})...")
self.simulation.step(protocol['step'])
#############
## WARNOUT ##
#############
state=self.simulation.context.getState(getPositions=True, getVelocities=True, getEnergy=True)
self.positions = state.getPositions()
self.velocities = state.getVelocities()
self.energy=state.getPotentialEnergy()
print('NVT equilibration finished.')
self.EQ_NPT_PDB=self.writePDB(self.simulation.topology, self.positions, name='equilibration_NVT')
self.system=self.forceHandler(self.system, kinds=['CustomExternalForce'])
#Remove implemented forces (only costum)
#TODO: remove MC for fresh delivery
return self.simulation
def test_reporter():
tempdir = os.path.join(dir, 'test1')
os.makedirs(tempdir)
pdb = PDBFile(get_fn('native.pdb'))
forcefield = ForceField('amber99sbildn.xml', 'amber99_obc.xml')
# NO PERIODIC BOUNARY CONDITIONS
system = forcefield.createSystem(pdb.topology,
nonbondedMethod=CutoffNonPeriodic,
nonbondedCutoff=1.0 * nanometers,
constraints=HBonds,
rigidWater=True)
integrator = LangevinIntegrator(300 * kelvin, 1.0 / picoseconds,
2.0 * femtoseconds)
integrator.setConstraintTolerance(0.00001)
platform = Platform.getPlatformByName('Reference')
simulation = Simulation(pdb.topology, system, integrator, platform)
simulation.context.setPositions(pdb.positions)
simulation.context.setVelocitiesToTemperature(300 * kelvin)
hdf5file = os.path.join(tempdir, 'traj.h5')
ncfile = os.path.join(tempdir, 'traj.nc')
dcdfile = os.path.join(tempdir, 'traj.dcd')
reporter = HDF5Reporter(hdf5file,
2,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True,
velocities=True)
reporter2 = NetCDFReporter(ncfile,
2,
coordinates=True,
time=True,
cell=True)
reporter3 = DCDReporter(dcdfile, 2)
simulation.reporters.append(reporter)
simulation.reporters.append(reporter2)
simulation.reporters.append(reporter3)
simulation.step(100)
reporter.close()
reporter2.close()
with HDF5TrajectoryFile(hdf5file) as f:
got = f.read()
yield lambda: eq(got.temperature.shape, (50, ))
yield lambda: eq(got.potentialEnergy.shape, (50, ))
yield lambda: eq(got.kineticEnergy.shape, (50, ))
yield lambda: eq(got.coordinates.shape, (50, 22, 3))
yield lambda: eq(got.velocities.shape, (50, 22, 3))
yield lambda: eq(got.cell_lengths, None)
yield lambda: eq(got.cell_angles, None)
yield lambda: eq(got.time, 0.002 * 2 * (1 + np.arange(50)))
yield lambda: f.topology == md.load(get_fn('native.pdb')).top
with NetCDFTrajectoryFile(ncfile) as f:
xyz, time, cell_lengths, cell_angles = f.read()
yield lambda: eq(cell_lengths, None)
yield lambda: eq(cell_angles, None)
yield lambda: eq(time, 0.002 * 2 * (1 + np.arange(50)))
hdf5_traj = md.load(hdf5file)
dcd_traj = md.load(dcdfile, top=get_fn('native.pdb'))
netcdf_traj = md.load(ncfile, top=get_fn('native.pdb'))
# we don't have to convert units here, because md.load already
# handles that
assert hdf5_traj.unitcell_vectors is None
yield lambda: eq(hdf5_traj.xyz, netcdf_traj.xyz)
yield lambda: eq(hdf5_traj.unitcell_vectors, netcdf_traj.unitcell_vectors)
yield lambda: eq(hdf5_traj.time, netcdf_traj.time)
yield lambda: eq(dcd_traj.xyz, hdf5_traj.xyz)
modeller = msm.convert(molsys, to_form='openmm.Modeller')
forcefield = app.ForceField("amber14-all.xml", "amber14/tip3p.xml")
system = forcefield.createSystem(modeller.topology,
nonbondedMethod=app.PME,
nonbondedCutoff=1.2 * unit.nanometer,
constraints=app.HBonds)
integrator = mm.LangevinIntegrator(300 * unit.kelvin, 1.0 / unit.picosecond,
2.0 * unit.femtoseconds)
platform = mm.Platform.getPlatformByName('CUDA')
simulation = app.Simulation(modeller.topology, system, integrator, platform)
simulation.context.setPositions(modeller.positions)
simulation.minimizeEnergy()
simulation.context.setVelocitiesToTemperature(300 * unit.kelvin)
simulation.reporters.append(
app.StateDataReporter(stdout,
50000,
progress=True,
potentialEnergy=True,
temperature=True,
remainingTime=True,
totalSteps=1000000))
simulation.reporters.append(
app.DCDReporter('villin_hp35_solvated.dcd', 50000,
enforcePeriodicBox=True))
simulation.reporters.append(HDF5Reporter('villin_hp35_solvated.h5', 50000))
simulation.step(1000000)
simulation.reporters[2].close()
final_positions = simulation.context.getState(getPositions=True).getPositions()
app.PDBFile.writeFile(modeller.topology, final_positions,
open('villin_hp35_solvated_last.pdb', 'w'))
speed=True,
step=True,
time=True,
potentialEnergy=True,
temperature=True,
volume=True,
totalSteps=steps_simulation,
separator=", "))
# Observables
simulation.reporters.append(
HDF5Reporter('trajectory.h5',
reportInterval=steps_interval_saving,
coordinates=True,
time=True,
cell=True,
potentialEnergy=True,
kineticEnergy=True,
temperature=True))
# Checkpoints
simulation.reporters.append(
app.CheckpointReporter('checkpnt.chk', steps_interval_checkpoint))
#### Running Simulation
start_simulation_realtime = realtime()
simulation.step(steps_simulation)
def buildSimulation(self,
integrator=LangevinMiddleIntegrator,
dt=0.002 * picoseconds,
temperature=298.15 * kelvin,
ensemble="NPT",
exceptions=[],
filePrefix="traj",
saveTrajectory=False,
trajInterval=500,
saveStateData=False,
stateDataInterval=250,
atomSubset=None,
thermalize=True):
"""
Build a simulation context from the system. The simulation is
then available as an attribute.
"""
if self.system is None:
raise AttributeError(
"Please first make a System using MDSystem.buildSystem()")
# If simulation exists, close any reporters
if self.simulation is not None:
for reporter in self.simulation.reporters:
try:
reporter.close()
except:
continue
# Setup MD simulation
integrator = integrator(temperature, 1 / picosecond, dt)
# Setup exceptions in nonbonded forces if provided
nonbonded = self.system.getForce(
self._getIndexOfNonbondedForce(self.system))
for atom1, atom2 in exceptions:
nonbonded.addException(int(atom1), int(atom2), 0.0, 0.0, 0.0, True)
# Setup barostat for NPT ensemble
if ensemble == "NPT":
barostat = MonteCarloBarostat(1.0 * bar, temperature, 25)
self.system.addForce(barostat)
# Add simulation
self.simulation = Simulation(self.topology, self.system, integrator)
# Initialize particle positions and velocities
self.simulation.context.setPositions(self.positions)
if thermalize:
self.thermalize(temperature=temperature)
# Add reporters
if saveTrajectory:
self.simulation.reporters.append(
HDF5Reporter(f"{filePrefix}.h5",
trajInterval,
atomSubset=atomSubset))
if saveStateData:
self.simulation.reporters.append(
StateDataReporter(f"{filePrefix}.csv",
stateDataInterval,
step=True,
time=True,
volume=True,
totalEnergy=True,
temperature=True,
elapsedTime=True))
return self