def __init__(self, bs_cutoff=0.8*unit.nanometer):
# must set this here since we need it to generate the state,
# will get called again in the superclass method
self.getState_kwargs = dict(GET_STATE_KWARG_DEFAULTS)
if self.GET_STATE_KWARGS is not None:
self.getState_kwargs.update(self.GET_STATE_KWARGS)
test_sys = LysozymeImplicit()
# set the box vectors to something reasonable
# we have a constructor which uses a constant we set in the
# class to do this
box_vectors = self.box_vectors()
# set them to the test system
test_sys.system.setDefaultPeriodicBoxVectors(
box_vectors[0],
box_vectors[1],
box_vectors[2],
)
init_state = self.make_state(test_sys.system, test_sys.positions)
lig_idxs = self.ligand_idxs()
bs_idxs = self.binding_site_idxs(bs_cutoff)
distance = UnbindingDistance(ligand_idxs=lig_idxs,
binding_site_idxs=bs_idxs,
ref_state=init_state)
super().__init__(
distance=distance,
init_state=init_state,
system=test_sys.system,
topology=test_sys.topology,
)
def main(n_runs,
n_cycles,
steps,
n_walkers,
n_workers=1,
debug_prints=False,
seed=None):
## Load objects needed for various purposes
# load a json string of the topology
with open(json_top_path, mode='r') as rf:
sEH_TPPU_system_top_json = rf.read()
# an openmm.State object for setting the initial walkers up
with open(omm_state_path, mode='rb') as rf:
omm_state = pickle.load(rf)
## set up the OpenMM Runner
# load the psf which is needed for making a system in OpenMM with
# CHARMM force fields
psf = omma.CharmmPsfFile(charmm_psf_path)
# set the box size lengths and angles
lengths = [CUBE_LENGTH for i in range(3)]
angles = [CUBE_ANGLE for i in range(3)]
psf.setBox(*lengths, *angles)
# charmm forcefields parameters
params = omma.CharmmParameterSet(*charmm_param_paths)
# create a system using the topology method giving it a topology and
# the method for calculation
system = psf.createSystem(params,
nonbondedMethod=omma.CutoffPeriodic,
nonbondedCutoff=NONBONDED_CUTOFF,
constraints=omma.HBonds)
# make this a constant temperature and pressure simulation at 1.0
# atm, 300 K, with volume move attempts every 50 steps
barostat = omm.MonteCarloBarostat(PRESSURE, TEMPERATURE, VOLUME_MOVE_FREQ)
# add it as a "Force" to the system
system.addForce(barostat)
# make an integrator object that is constant temperature
integrator = omm.LangevinIntegrator(TEMPERATURE, FRICTION_COEFFICIENT,
STEP_SIZE)
# set up the OpenMMRunner with the system
runner = OpenMMRunner(system, psf.topology, integrator, platform=PLATFORM)
# the initial state, which is used as reference for many things
init_state = OpenMMState(omm_state)
## Make the distance Metric
# load the crystal structure coordinates
crystal_traj = mdj.load_pdb(pdb_path)
# get the atoms in the binding site according to the crystal structure
bs_idxs = binding_site_atoms(crystal_traj.top, LIG_RESID,
crystal_traj.xyz[0])
lig_idxs = ligand_idxs(crystal_traj.top, LIG_RESID)
prot_idxs = protein_idxs(crystal_traj.top)
# make the distance metric with the ligand and binding site
# indices for selecting atoms for the image and for doing the
# alignments to only the binding site. All images will be aligned
# to the reference initial state
unb_distance = UnbindingDistance(lig_idxs, bs_idxs, init_state)
## Make the resampler
# make a Wexplore resampler with default parameters and our
# distance metric
resampler = WExploreResampler(distance=unb_distance,
init_state=init_state,
max_n_regions=MAX_N_REGIONS,
max_region_sizes=MAX_REGION_SIZES,
pmin=PMIN,
pmax=PMAX)
## Make the Boundary Conditions
# makes ref_traj and selects lingand_atom and protein atom indices
# instantiate a revo unbindingboudaryconditiobs
ubc = UnbindingBC(cutoff_distance=CUTOFF_DISTANCE,
initial_state=init_state,
topology=crystal_traj.topology,
ligand_idxs=lig_idxs,
receptor_idxs=prot_idxs)
## make the reporters
# WepyHDF5
# make a dictionary of units for adding to the HDF5
# open it in truncate mode first, then switch after first run
hdf5_reporter = WepyHDF5Reporter(
hdf5_path,
mode='w',
# the fields of the State that will be saved in the HDF5 file
save_fields=SAVE_FIELDS,
# the topology in a JSON format
topology=sEH_TPPU_system_top_json,
# the resampler and boundary
# conditions for getting data
# types and shapes for saving
resampler=resampler,
boundary_conditions=ubc,
# the units to save the fields in
units=dict(UNITS),
# sparse (in time) fields
sparse_fields=dict(SPARSE_FIELDS),
# sparse atoms fields
main_rep_idxs=np.concatenate((lig_idxs, prot_idxs)),
all_atoms_rep_freq=ALL_ATOMS_SAVE_FREQ)
dashboard_reporter = WExploreDashboardReporter(
dashboard_path,
mode='w',
step_time=STEP_SIZE.value_in_unit(unit.second),
max_n_regions=resampler.max_n_regions,
max_region_sizes=resampler.max_region_sizes,
bc_cutoff_distance=ubc.cutoff_distance)
setup_reporter = SetupReporter(setup_state_path, mode='w')
restart_reporter = RestartReporter(restart_state_path, mode='w')
reporters = [
hdf5_reporter, dashboard_reporter, setup_reporter, restart_reporter
]
## The work mapper
# we use a mapper that uses GPUs
work_mapper = WorkerMapper(worker_type=OpenMMGPUWorker,
num_workers=n_workers)
## Combine all these parts and setup the simulation manager
# set up parameters for running the simulation
# initial weights
init_weight = 1.0 / n_walkers
# a list of the initial walkers
init_walkers = [
Walker(OpenMMState(omm_state), init_weight) for i in range(n_walkers)
]
# Instantiate a simulation manager
sim_manager = Manager(init_walkers,
runner=runner,
resampler=resampler,
boundary_conditions=ubc,
work_mapper=work_mapper,
reporters=reporters)
### RUN the simulation
for run_idx in range(n_runs):
print("Starting run: {}".format(run_idx))
sim_manager.run_simulation(n_cycles, steps, debug_prints=True)
print("Finished run: {}".format(run_idx))
def run_sim(init_state_path,
json_top_path,
forcefield_paths,
n_cycles,
n_steps,
platform,
n_workers,
lig_ff=None,
**kwargs):
# add in the ligand force fields
assert lig_ff is not None, "must give ligand forcefield"
forcefield_paths.append(lig_ff)
#### Wepy Orchestrator
# load the wepy.OpenMMState
with open(init_state_path, 'rb') as rf:
init_state = pickle.load(rf)
### Apparatus
# Runner components
# load the JSON for the topology
with open(json_top_path) as rf:
json_top_str = rf.read()
# load it with mdtraj and then convert to openmm
mdj_top = json_to_mdtraj_topology(json_top_str)
omm_topology = mdj_top.to_openmm()
# we need to use the box vectors for setting the simulation up,
# paying mind to the units
box_vectors = init_state['box_vectors'] * init_state.box_vectors_unit
positions = init_state['positions'] * init_state.positions_unit
# set the box to the last box size from equilibration
omm_topology.setPeriodicBoxVectors(box_vectors)
# force field parameters
force_field = omma.ForceField(*forcefield_paths)
# create a system using the topology method giving it a topology and
# the method for calculation
runner_system = force_field.createSystem(omm_topology,
nonbondedMethod=NONBONDED_METHOD,
nonbondedCutoff=NONBONDED_CUTOFF,
constraints=MD_CONSTRAINTS,
rigidWater=RIGID_WATER,
removeCMMotion=REMOVE_CM_MOTION,
hydrogenMass=HYDROGEN_MASS)
# barostat to keep pressure constant
runner_barostat = omm.MonteCarloBarostat(PRESSURE, TEMPERATURE,
VOLUME_MOVE_FREQ)
# add it to the system
runner_system.addForce(runner_barostat)
# set up for a short simulation to runner and prepare
# instantiate an integrator
runner_integrator = omm.LangevinIntegrator(TEMPERATURE,
FRICTION_COEFFICIENT, STEP_TIME)
## Runner
runner = OpenMMRunner(runner_system,
omm_topology,
runner_integrator,
platform=platform)
## Resampler
# Distance Metric
lig_idxs = ligand_idxs(json_top_str)
prot_idxs = protein_idxs(json_top_str)
bs_idxs = binding_site_idxs(json_top_str, positions, box_vectors, CUTOFF)
# set distance metric
distance_metric = UnbindingDistance(lig_idxs, bs_idxs, init_state)
# set resampler
resampler = WExploreResampler(distance=distance_metric,
init_state=init_state,
max_n_regions=MAX_N_REGIONS,
max_region_sizes=MAX_REGION_SIZES,
pmin=PMIN,
pmax=PMAX)
## Boundary Conditions
# optional: set the boundary conditions
bc = None
## CONFIGURATION
# the idxs of the main representation to save in the output files,
# it is just the protein and the ligand
# optional: set the main representation atom indices, set to None
# to save all the atoms in the 'positions' field
main_rep_idxs = np.concatenate((lig_idxs, prot_idxs))
# REPORTERS
# list of reporter classes and partial kwargs for using in the
# orchestrator
hdf5_reporter_kwargs = {
'main_rep_idxs': main_rep_idxs,
'topology': json_top_str,
'resampler': resampler,
'boundary_conditions': bc,
# general parameters
'save_fields': SAVE_FIELDS,
'units': dict(UNITS),
'sparse_fields': dict(SPARSE_FIELDS),
'all_atoms_rep_freq': ALL_ATOMS_SAVE_FREQ
}
# get all the reporters together. Order is important since they
# will get paired with the kwargs
reporter_classes = [
WepyHDF5Reporter,
]
# collate the kwargs in the same order
reporter_kwargs = [
hdf5_reporter_kwargs,
]
# make the configuration with all these reporters and the default
# number of workers. Don't be thrown off by this. You don't need
# this. It is just a convenient way to dynamically name the
# outputs of the reporters and parametrize the workers and worker
# mappers. This is mainly for use in the Orchestrator framework
# but it is useful here just for batch naming everything.
configuration = Configuration(n_workers=DEFAULT_N_WORKERS,
reporter_classes=reporter_classes,
reporter_partial_kwargs=reporter_kwargs,
config_name="no-orch",
mode='w')
# then instantiate the reporters from the configuration. THis
# localizes the file paths to outputs and applies the key-word
# arguments specified above.
reporters = configuration._gen_reporters()
print("created configuration")
### Initial Walkers
init_walkers = [
Walker(deepcopy(init_state), INIT_WEIGHT) for _ in range(N_WALKERS)
]
print("created init walkers")
### Work Mapper
if platform in ('OpenCL', 'CUDA'):
# we use a mapper that uses GPUs
work_mapper = WorkerMapper(worker_type=OpenMMGPUWorker,
num_workers=n_workers)
elif platform in ('CPU', ):
# for the CPU we can choose how many threads to use per walker.
worker_attributes = {'num_threads': N_CPU_THREADS}
work_mapper = WorkerMapper(worker_type=OpenMMCPUWorker,
worker_attributes=worker_attributes,
num_workers=n_workers)
elif platform in ('Reference', ):
# we just use the standard mapper for in serial
work_mapper = Mapper
### Simulation Manager
sim_manager = Manager(init_walkers,
runner=runner,
resampler=resampler,
boundary_conditions=bc,
work_mapper=work_mapper,
reporters=reporters)
### Run the simulation
steps = [n_steps for _ in range(n_cycles)]
sim_manager.run_simulation(n_cycles, steps)
## Runner
RUNNER = OpenMMRunner(system, psf.topology, integrator, platform=PLATFORM)
# the initial state, which is used as reference for many things
INIT_STATE = OpenMMState(MINIMIZED_INIT_OMM_STATE)
## Resampler
# Distance Metric
# make the distance metric with the ligand and binding site
# indices for selecting atoms for the image and for doing the
# alignments to only the binding site. All images will be aligned
# to the reference initial state
DISTANCE_METRIC = UnbindingDistance(LIG_IDXS, BS_IDXS, INIT_STATE)
# WExplore Resampler
# make a Wexplore resampler with default parameters and our
# distance metric
RESAMPLER = WExploreResampler(distance=DISTANCE_METRIC,
init_state=INIT_STATE,
max_n_regions=MAX_N_REGIONS,
max_region_sizes=MAX_REGION_SIZES,
pmin=PMIN, pmax=PMAX)
## Boundary Conditions
# makes ref_traj and selects lingand_atom and protein atom indices
# instantiate a revo unbindingboudaryconditiobs