def walkers_from_disk(n_expected_walkers=48, path='$WEST_SIM_ROOT/traj_segs/'):
# find the last iteration that was completed; this is for emergency cases and other hacks
path = os.path.expandvars(path)
max_iteration = -1
n_found_walkers = {}
for topdir in [d for d in _get_dirs(path) if d.isdigit()]:
iteration = int(topdir)
subdirs = [d for d in _get_dirs(path + topdir) if d.isdigit()]
if len(
set([int(d) for d in subdirs])
& set(range(n_expected_walkers))) == n_expected_walkers:
max_iteration = max(max_iteration, iteration)
n_found_walkers[iteration] = len(set([int(d) for d in subdirs]))
if max_iteration == -1:
raise RuntimeError('no valid iteration found')
weights = np.ones(n_expected_walkers, dtype=float) / n_expected_walkers
assert n_found_walkers[max_iteration] == n_expected_walkers
walkers = [
Walker(WestpaWalkerState.from_file(iteration=max_iteration, id=i),
weight=weights[i])
for i in range(n_found_walkers[max_iteration])
]
print('continuing at iteration', max_iteration, 'with',
n_found_walkers[max_iteration], 'walkers and with discarded weights')
return walkers, max_iteration
def task_fail(walker):
n = walker.state['num']
if n == 1:
raise ValueError("No soup for you!!")
else:
return Walker(WalkerState(**{'num' : n+1}), walker.weight)
def make_initial_walkers(cls, state, n_walkers):
init_weight = 1.0 / n_walkers
init_walkers = [
Walker(deepcopy(state), init_weight) for i in range(n_walkers)
]
return init_walkers
def _runtest(self, num_walkers, num_cycles, dimension, debug_prints=False):
print("Random walk simulation with: ")
print("Dimension =", dimension)
print("Probability =", self.probability)
print("Number of Walkers", num_walkers)
print("Number of Cycles", num_cycles)
# set up initial state for walkers
positions = np.zeros((1, dimension))
init_state = WalkerState(positions=positions, time=0.0)
# create list of init_walkers
initial_weight = 1 / num_walkers
init_walkers = []
# init_walkers, n_cycles = get_final_state(path, num_walkers)
init_walkers = [
Walker(init_state, initial_weight) for i in range(num_walkers)
]
# set up raunner for system
runner = RandomWalkRunner(dimension=dimension,
probability=self.probability)
units = dict(UNIT_NAMES)
# instantiate a revo unbindingboudaryconditiobs
segment_length = 10
# set up the reporter
randomwalk_system_top_json = self.generate_topology()
hdf5_reporter = WepyHDF5Reporter(self.hdf5_reporter_path,
mode='w',
save_fields=SAVE_FIELDS,
topology=randomwalk_system_top_json,
resampler=self.resampler,
units=dict(UNITS),
n_dims=dimension)
# running the simulation
sim_manager = Manager(init_walkers,
runner=runner,
resampler=self.resampler,
work_mapper=Mapper(),
reporters=[hdf5_reporter])
# run a simulation with the manager for n_steps cycles of length 1000 each
steps = [segment_length for i in range(num_cycles)]
print("Start simulation")
sim_manager.run_simulation(num_cycles,
steps,
debug_prints=debug_prints)
print("Finished Simulation")
def run_segment(self,
walker,
segment_length,
propagation_id,
random_seeds,
debug_prints=False):
# segment_length is ignored for now
parent_id = walker.state.id
iteration = walker.state[
'iteration'] + 1 # TODO: increment here or later?
root = self.west_sim_root
env = dict(os.environ)
if parent_id == -1: # start new trajectory
assert walker.state.struct_data_ref is not None
env['WEST_PARENT_DATA_REF'] = walker.state.struct_data_ref
env['WEST_CURRENT_SEG_INITPOINT_TYPE'] = 'SEG_INITPOINT_NEWTRAJ'
else: # continue trajectory
env['WEST_PARENT_DATA_REF'] = '%s/traj_segs/%06d/%06d' % (
root, iteration - 1, parent_id)
env['WEST_CURRENT_SEG_INITPOINT_TYPE'] = 'SEG_INITPOINT_CONTINUES'
env['WEST_CURRENT_ITER'] = '%d' % iteration
env['WEST_CURRENT_SEG_ID'] = '%d' % propagation_id
env['WEST_PARENT_ID'] = '%d' % parent_id
env['WEST_CURRENT_SEG_DATA_REF'] = '%s/traj_segs/%06d/%06d' % (
root, iteration, propagation_id)
env.update(random_seeds)
pcoor_file = tempfile.NamedTemporaryFile(mode='r')
env['WEST_PCOORD_RETURN'] = pcoor_file.name
env['WEST_COORD_RETURN'] = '/dev/null'
env['SEG_DEBUG'] = ''
erl = subprocess.call(os.path.expandvars(self.runseg),
shell=True,
env=env)
if erl != 0:
raise RuntimeError('segment propagation failed')
# creates new_walker from new state and current weight
pcoor = np.loadtxt(pcoor_file)
pcoor_file.close()
new_state = WestpaWalkerState(positions=np.atleast_2d(pcoor)[-1, :],
iteration=iteration,
parent_id=parent_id,
id=propagation_id)
if debug_prints:
print('walker #', id, 'with parent', parent_id, 'has weight',
walker.weight)
new_walker = Walker(state=new_state, weight=walker.weight)
return new_walker
def main(n_walkers=36,
n_workers=12,
n_runs=1,
n_cycles=20,
n_steps=100,
continue_sim=False):
runner = WestpaRunner()
init_state = WestpaWalkerState.from_bstate(
struct_data_ref='$WEST_SIM_ROOT/bstates/0')
work_mapper = WorkerMapper(worker_type=Worker, num_workers=n_workers)
if continue_sim:
# init_walkers = walkers_from_disk(n_expected_walkers=n_walkers)
init_walkers, start_cycle = WalkersPickleReporter.load_most_recent_cycle(
debug_prints=True)
else:
start_cycle = 0
init_weight = 1.0 / n_walkers
init_walkers = [
Walker(deepcopy(init_state), init_weight) for i in range(n_walkers)
]
unb_distance = PairDistance()
resampler = REVOResampler(distance=unb_distance, init_state=init_state)
reporters = [WalkersPickleReporter(freq=10)]
# Instantiate a simulation manager
sim_manager = Manager(init_walkers,
runner=runner,
resampler=resampler,
boundary_conditions=NoBC(),
work_mapper=work_mapper,
reporters=reporters)
segment_lengths = [n_steps for i in range(start_cycle + n_cycles)]
### RUN the simulation
for run_idx in range(n_runs):
print("Starting run: {}".format(run_idx))
sim_manager.continue_run_simulation(
0,
n_cycles,
segment_lengths,
num_workers=n_workers,
start_cycle=start_cycle) #, debug_prints=True)
print("Finished run: {}".format(run_idx))
def run_segment(self, walker, segment_length):
# documented in superclass
# Gets the current posiotion of RandomWalk Walker
positions = walker.state['positions']
# Make movements for the segment_length steps
for _ in range(segment_length):
# calls walk function for one step movement
new_positions = self._walk(positions)
positions = new_positions
# makes new state form new positions
new_state = WalkerState(positions=new_positions, time=0.0)
# creates new_walker from new state and current weight
new_walker = Walker(new_state, walker.weight)
return new_walker
def run_segment(self, walker, segment_length, **kwargs):
"""Runs a random walk simulation for the given number of steps.
Parameters
----------
walker : object implementing the Walker interface
The walker for which dynamics will be propagated.
segment_length : int
The numerical value that specifies how much dynamical steps
are to be run.
Returns
-------
new_walker : object implementing the Walker interface
Walker after dynamics was run, only the state should be modified.
"""
# Gets the current posiotion of RandomWalk Walker
positions = walker.state['positions']
# Make movements for the segment_length steps
for _ in range(segment_length):
# calls walk function for one step movement
new_positions = self._walk(positions)
positions = new_positions
# makes new state form new positions
new_state = WalkerState(positions=new_positions, time=0.0)
# creates new_walker from new state and current weight
new_walker = Walker(new_state, walker.weight)
return new_walker
mapper = Mapper()
## Run the simulation
if __name__ == "__main__":
# normalize the output paths
hdf5_path = osp.join(outputs_dir, hdf5_filename)
print("Number of steps: {}".format(n_steps))
print("Number of cycles: {}".format(n_cycles))
# # create the initial walkers
init_weight = 1.0 / n_walkers
init_walkers = [
Walker(OpenMMState(omm_states[i], activity=np.array([0.])),
init_weight) for i in range(n_walkers)
]
hdf5_reporter = WepyHDF5Reporter(
file_path=hdf5_path,
mode='w',
# save_fields set to None saves everything
save_fields=None,
resampler=resampler,
boundary_conditions=None,
topology=json_top,
units=units)
reporters = [hdf5_reporter]
sim_manager = Manager(init_walkers,
runner=runner,
if __name__ == "__main__":
if sys.argv[1] == "-h" or sys.argv[1] == "--help":
print("arguments: n_cycles, n_steps, n_walkers")
else:
n_cycles = int(sys.argv[1])
n_steps = int(sys.argv[2])
n_walkers = int(sys.argv[3])
print("Number of steps: {}".format(n_steps))
print("Number of cycles: {}".format(n_cycles))
# create the initial walkers
init_weight = 1.0 / n_walkers
init_walkers = [Walker(OpenMMState(init_sim_state), init_weight) for i in range(n_walkers)]
# initialize the simulation manager
sim_manager = Manager(init_walkers,
runner=runner,
resampler=resampler,
boundary_conditions=ubc,
work_mapper=mapper,
reporters=reporters)
# make a number of steps for each cycle. In principle it could be
# different each cycle
steps = [n_steps for i in range(n_cycles)]
# actually run the simulation
print("Starting run: {}".format(0))
def task_pass(walker):
# simulate it actually taking some time
n = walker.state['num']
return Walker(WalkerState(**{'num' : n+1}), walker.weight)
if __name__ == "__main__":
if sys.argv[1] == "-h" or sys.argv[1] == "--help":
print("arguments: n_cycles, n_steps, n_walkers")
else:
n_cycles = int(sys.argv[1])
n_steps = int(sys.argv[2])
n_walkers = int(sys.argv[3])
print("Number of steps: {}".format(n_steps))
print("Number of cycles: {}".format(n_cycles))
# create the initial walkers
init_weight = 1.0 / n_walkers
init_walkers = [
Walker(OpenMMState(init_sim_state), init_weight)
for i in range(n_walkers)
]
# initialize the simulation manager
sim_manager = Manager(init_walkers,
runner=runner,
resampler=resampler,
boundary_conditions=ubc,
work_mapper=mapper,
reporters=reporters)
# make a number of steps for each cycle. In principle it could be
# different each cycle
steps = [n_steps for i in range(n_cycles)]
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 get_char_distance(dimension, num_walkers):
"""Calculate the characteristic value.
Runs one cycle simulation and calculates the characteristic
distance value.
Parameters
----------
dimension: int
The dimension of the random walk space.
num_walkers: int
The number of walkers.
Returns
-------
characteristic distance : float
The characteristic distance value.
"""
# set up initial state for walkers
positions = np.zeros((1, dimension))
init_state = WalkerState(positions=positions, time=0.0)
# set up the distance function
rw_distance = RandomWalkDistance()
# set up the REVO Resampler with the parameters
resampler = REVOResampler(distance=rw_distance,
pmin=PMIN,
pmax=PMAX,
init_state=init_state,
char_dist=1,
merge_dist=MERGE_DIST)
# create list of init_walkers
initial_weight = 1 / num_walkers
init_walkers = [
Walker(init_state, initial_weight) for i in range(num_walkers)
]
# set up raunner for system
runner = RandomWalkRunner(probability=PROBABILITY)
n_steps = 10
mapper = Mapper()
# running the simulation
sim_manager = Manager(init_walkers,
runner=runner,
resampler=resampler,
work_mapper=mapper)
print("Running simulation")
#runs for one cycle
sim_manager.init(num_walkers)
new_walkers = sim_manager.run_segment(init_walkers, n_steps, 0)
dist_matrix, _ = resampler._all_to_all_distance(new_walkers)
return np.average(dist_matrix)
sim_apparatus = WepySimApparatus(runner,
resampler=resampler,
boundary_conditions=ubc)
# we also create a default configuration for the orchestrator that
# will be used unless one is given at runtime for the creation of a
# simulation manager
configuration = Configuration(n_workers=4,
reporter_classes=reporter_classes,
reporter_partial_kwargs=reporter_kwargs)
# we also want to set up the orchestrator with some default walkers to
# use to get us started. Otherwise these could be provided from a
# snapshot or on their own. Ideally we only want to have a single
# script setting up an orchestrator and then manage everything else on
# the command line intereactively from then on out
init_weight = 1.0 / N_WALKERS
init_walkers = [
Walker(OpenMMState(init_sim_state), init_weight) for i in range(N_WALKERS)
]
# then create the seed/root/master orchestrator which will be used
# from here on out
orchestrator = Orchestrator(sim_apparatus,
default_init_walkers=init_walkers,
default_configuration=configuration)
# save it
orchestrator.dump('LJ-pair.orch', mode='wb')
def init(self, continue_run=None, init_walkers=None, **kwargs):
# do the inherited stuff
super().init(**kwargs)
# open and initialize the HDF5 file
logging.info("Initializing HDF5 file at {}".format(self.file_path))
self.wepy_h5 = WepyHDF5(self.file_path,
mode=self.mode,
topology=self._tmp_topology,
units=self.units,
sparse_fields=list(self._sparse_fields.keys()),
feature_shapes=self._feature_shapes,
feature_dtypes=self._feature_dtypes,
n_dims=self._n_dims,
main_rep_idxs=self.main_rep_idxs,
alt_reps=self.alt_reps_idxs)
# if we specify save fields only save these for the initial walkers
if self.save_fields is not None:
state_fields = list(init_walkers[0].state.dict().keys())
# make sure all the save_fields are present in the state
assert all([True if save_field in state_fields else False
for save_field in self.save_fields]), \
"Not all specified save_fields present in walker states"
filtered_init_walkers = []
for walker in init_walkers:
# make a new state by filtering the attributes of the old ones
state_d = {
k: v
for k, v in walker.state.dict().items()
if k in self.save_fields
}
# and saving alternate representations as we would
# expect them
# if there are any alternate representations set them
for alt_rep_name, alt_rep_idxs in self.alt_reps_idxs.items():
alt_rep_path = 'alt_reps/{}'.format(alt_rep_name)
# if the idxs are None we want all of the atoms
if alt_rep_idxs is None:
state_d[alt_rep_path] = state_d['positions'][:]
# otherwise get only the atoms we want
else:
state_d[alt_rep_path] = state_d['positions'][
alt_rep_idxs]
# if the main rep is different then the full state
# positions set that
if self.main_rep_idxs is not None:
state_d['positions'] = state_d['positions'][
self.main_rep_idxs]
# then making the new state
new_state = WalkerState(**state_d)
filtered_init_walkers.append(Walker(new_state, walker.weight))
# otherwise save the full state
else:
filtered_init_walkers = init_walkers
self.wepy_h5.set_mode(mode='r+')
with self.wepy_h5:
# if this is a continuation run of another run we want to
# initialize it as such
# initialize a new run
run_grp = self.wepy_h5.new_run(filtered_init_walkers,
continue_run=continue_run)
self.wepy_run_idx = run_grp.attrs['run_idx']
# initialize the run record groups using their fields
self.wepy_h5.init_run_fields_resampling(self.wepy_run_idx,
self.resampling_fields)
# the enumeration for the values of resampling
self.wepy_h5.init_run_fields_resampling_decision(
self.wepy_run_idx, self.decision_enum)
self.wepy_h5.init_run_fields_resampler(self.wepy_run_idx,
self.resampler_fields)
# set the fields that are records for tables etc. unless
# they are already set
if 'resampling' not in self.wepy_h5.record_fields:
self.wepy_h5.init_record_fields('resampling',
self.resampling_records)
if 'resampler' not in self.wepy_h5.record_fields:
self.wepy_h5.init_record_fields('resampler',
self.resampler_records)
# if there were no warping fields set there is no boundary
# conditions and we don't initialize them
if self.warping_fields is not None:
self.wepy_h5.init_run_fields_warping(self.wepy_run_idx,
self.warping_fields)
self.wepy_h5.init_run_fields_progress(self.wepy_run_idx,
self.progress_fields)
self.wepy_h5.init_run_fields_bc(self.wepy_run_idx,
self.bc_fields)
# table records
if 'warping' not in self.wepy_h5.record_fields:
self.wepy_h5.init_record_fields('warping',
self.warping_records)
if 'boundary_conditions' not in self.wepy_h5.record_fields:
self.wepy_h5.init_record_fields('boundary_conditions',
self.bc_records)
if 'progress' not in self.wepy_h5.record_fields:
self.wepy_h5.init_record_fields('progress',
self.progress_records)
# if this was opened in a truncation mode, we don't want to
# overwrite old runs with future calls to init(). so we
# change the mode to read/write 'r+'
if self.mode == 'w':
self.set_mode(0, 'r+')
n_cycles = 1
# the number of MD dynamics steps for each cycle
n_steps = 1000000
steps = [n_steps for i in range(n_cycles)]
# number of parallel simulations
n_walkers = 10
# the work mapper
# work_mapper = ThreadMapper()
work_mapper = Mapper()
# create the initial walkers with equal weights
with start_action(action_type="Init Walkers") as ctx:
init_weight = 1.0 / n_walkers
init_walkers = [Walker(copy(init_state), init_weight) for i in range(n_walkers)]
with start_action(action_type="Init Sim Manager") as ctx:
sim_manager = Manager(
init_walkers,
runner=runner,
resampler=resampler,
work_mapper=work_mapper)
# run the simulation and get the results
with start_action(action_type="Simulation") as ctx:
final_walkers, _ = sim_manager.run_simulation(n_cycles, steps)
'max_n_regions' : RESAMPLER.max_n_regions,
'max_region_sizes' : RESAMPLER.max_region_sizes,
'bc_cutoff_distance' : BC.cutoff_distance}
REPORTER_KWARGS = [hdf5_reporter_kwargs, dashboard_reporter_kwargs]
N_WORKERS = 8
CONFIGURATION = Configuration(n_workers=N_WORKERS,
reporter_classes=REPORTER_CLASSES,
reporter_partial_kwargs=REPORTER_KWARGS)
print("created configuration")
### Initial Walkers
N_WALKERS = 48
INIT_WEIGHT = 1.0 / N_WALKERS
INIT_WALKERS = [Walker(deepcopy(INIT_STATE), INIT_WEIGHT) for i in range(N_WALKERS)]
print("created init walkers")
### Orchestrator
ORCHESTRATOR = Orchestrator(APPARATUS,
default_init_walkers=INIT_WALKERS,
default_configuration=CONFIGURATION)
print("created orchestrator, creating object file now")
ORCH_NAME = "sEH-TPPU"
dump_orchestrator(ORCHESTRATOR, "{}.orch".format(ORCH_NAME), mode='wb')
def gen_walkers(n_args):
args = range(n_args)
return [Walker(WalkerState(**{'num': arg}), 1 / len(args)) for arg in args]
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)
def gen_walkers():
return [Walker(WalkerState(**{'num' : arg}), 1/len(ARGS))
for arg in ARGS]
def test_apparatus_configuration(datadir_factory, mocker):
config = Configuration()
assert config.apparatus_opts == {}
reparam_config = config.reparametrize(apparatus_opts={
'runner': {
'platform': 'CPU',
},
})
assert reparam_config.apparatus_opts == {
'runner': {
'platform': 'CPU',
},
}
## test that we can change the apparatus parameters in the sim_manager
system_mock = mocker.Mock()
topology_mock = mocker.Mock()
integrator_mock = mocker.Mock()
runner = OpenMMRunner(
system_mock,
topology_mock,
integrator_mock,
)
resampler_mock = mocker.MagicMock() # mocker.patch('WExploreResampler')
apparatus = WepySimApparatus(
runner,
resampler=resampler_mock,
boundary_conditions=None,
)
apparatus._filters = (
runner,
None,
resampler_mock,
)
state_mock = mocker.MagicMock(
) # mocker.patch('wepy.walker.WalkerState', autospec=True)
walkers = [Walker(state_mock, 0.1) for i in range(1)]
snapshot = SimSnapshot(
walkers,
apparatus,
)
snapshot._walkers = walkers
snapshot._apparatus = apparatus
datadir = datadir_factory.mkdatadir()
orch = Orchestrator(orch_path=str(datadir / "test.orch.sqlite3"))
sim_manager = orch.gen_sim_manager(
snapshot,
reparam_config,
)
sim_manager.init()
# sim_mock = mocker.patch('wepy.runners.openmm.omma.Simulation')
# platform_mock = mocker.patch('wepy.runners.openmm.omm.Platform')
# _ = sim_manager.run_cycle(
# walkers,
# 2,
# 0,
# runner_opts={
# 'platform' : 'CPU',
# }
# )
# platform_mock.getPlatformByName.assert_called_with('CPU')
sim_mock = mocker.patch('wepy.runners.openmm.omma.Simulation')
platform_mock = mocker.patch('wepy.runners.openmm.omm.Platform')
platform_mock.getPlatformByName.\
return_value.getPropertyNames.\
return_value = ('Threads',)
_ = sim_manager.run_cycle(walkers,
2,
0,
runner_opts={
'platform': 'CPU',
'platform_kwargs': {
'Threads': '3'
},
})
platform_mock.getPlatformByName.assert_called_with('CPU')
platform_mock.getPlatformByName.\
return_value.getPropertyNames.\
assert_called()
platform_mock.getPlatformByName.\
return_value.setPropertyDefaultValue.\
assert_called_with(
'Threads',
'3'
)
def run_sim(init_state_path, json_top_path, forcefield_paths, n_cycles,
n_steps, n_workers, **kwargs):
#### 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
# 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
# TODO set distance metric
distance_metric = None
# TODO set resampler
resampler = None
## Boundary Conditions
# TODO optional: set the boundary conditions
bc = None
# apparatus = WepySimApparatus(runner, resampler=resampler,
# boundary_conditions=bc)
print("created apparatus")
## CONFIGURATION
# the idxs of the main representation to save in the output files,
# it is just the protein and the ligand
# TODO optional: set the main representation atom indices
main_rep_idxs = None
# 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
configuration = Configuration(n_workers=DEFAULT_N_WORKERS,
reporter_classes=reporter_classes,
reporter_partial_kwargs=reporter_kwargs,
config_name="no-orch")
# then instantiate them
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")
### Orchestrator
# orchestrator = Orchestrator(apparatus,
# default_init_walkers=init_walkers,
# default_configuration=configuration)
### Work Mapper
if PLATFORM in ('OpenCL', 'CUDA'):
# we use a mapper that uses GPUs
work_mapper = WorkerMapper(worker_type=OpenMMGPUWorker,
num_workers=n_workers)
if PLATFORM in ('Reference', 'CPU'):
# we just use the standard mapper
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)
def _run(self, num_runs, num_cycles, num_walkers):
"""Runs a random walk simulation.
Parameters
----------
num_runs: int
The number independet simulations.
num_cycles: int
The number of cycles that will be run in the simulation.
num_walkers: int
The number of walkers.
"""
print("Random walk simulation with: ")
print("Dimension = {}".format(self.dimension))
print("Probability = {}".format(self.probability))
print("Number of Walkers = {}".format(num_walkers))
print("Number of Cycles ={}".format(num_cycles))
# set up initial state for walkers
positions = np.zeros((1, self.dimension))
init_state = WalkerState(positions=positions, time=0.0)
# create list of init_walkers
initial_weight = 1 / num_walkers
init_walkers = []
init_walkers = [
Walker(init_state, initial_weight) for i in range(num_walkers)
]
# set up raunner for system
runner = RandomWalkRunner(probability=self.probability)
units = dict(UNIT_NAMES)
# instantiate a revo unbindingboudaryconditiobs
segment_length = 10
# set up the reporter
randomwalk_system_top_json = self.generate_topology()
hdf5_reporter = WepyHDF5Reporter(file_path=self.hdf5_filename,
mode='w',
save_fields=SAVE_FIELDS,
topology=randomwalk_system_top_json,
resampler=self.resampler,
units=dict(UNITS),
n_dims=self.dimension)
# running the simulation
sim_manager = Manager(init_walkers,
runner=runner,
resampler=self.resampler,
work_mapper=Mapper(),
reporters=[hdf5_reporter])
# run a simulation with the manager for n_steps cycles of length 1000 each
steps = [segment_length for i in range(num_cycles)]
### RUN the simulation
for run_idx in range(num_runs):
print("Starting run: {}".format(run_idx))
sim_manager.run_simulation(num_cycles, steps)
print("Finished run: {}".format(run_idx))
print("Finished Simulation")