def lj_openmm_runner(request, lj_omm_sys, lj_integrator):
# parametrize the platform
platform = request.param
positions = test_sys.positions.value_in_unit(test_sys.positions.unit)
init_state = gen_walker_state(positions, test_sys.system, integrator)
# initialize the runner
runner = OpenMMRunner(lj_omm_sys.system,
lj_omm_sys.topology,
lj_integrator,
platform=platform)
return runner
integrator = omm.LangevinIntegrator(TEMPERATURE, FRICTION_COEFFICIENT,
STEP_SIZE)
# make a context and set the positions
context = omm.Context(test_sys.system, copy(integrator))
context.setPositions(test_sys.positions)
# get the data from this context so we have a state to start the
# simulation with
get_state_kwargs = dict(GET_STATE_KWARG_DEFAULTS)
init_sim_state = context.getState(**get_state_kwargs)
init_state = OpenMMState(init_sim_state)
# initialize the runner
runner = OpenMMRunner(test_sys.system,
test_sys.topology,
integrator,
platform=PLATFORM)
## Distance Metric
# we define a simple distance metric for this system, assuming the
# positions are in a 'positions' field
class PairDistance(Distance):
def __init__(self, metric=euclidean):
self.metric = metric
def image(self, state):
return state['positions']
def image_distance(self, image_a, image_b):
dist_a = self.metric(image_a[0], image_a[1])
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))
from wepy.work_mapper.mapper import Mapper
# from wepy.work_mapper.thread import ThreadMapper
# use a ready made system for OpenMM MD simulation
with start_action(action_type="Instantiate test system"):
test_sys = LennardJonesPair()
with start_action(action_type="Gen Runner"):
integrator = omm.LangevinIntegrator(300.0*unit.kelvin,
1/unit.picosecond,
0.002*unit.picoseconds)
init_state = gen_walker_state(test_sys.positions, test_sys.system, integrator)
runner = OpenMMRunner(test_sys.system, test_sys.topology, integrator,
platform='Reference')
# a trivial resampler which does nothing
with start_action(action_type="Instantiate Resampler"):
resampler = NoResampler()
# Run the simulation
# number of cycles of WE to perform
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
def make_apparatus(
self,
platform='Reference',
platform_params=None,
runner_params=None,
integrator='LangevinIntegrator',
integrator_params=None,
resampler='WExploreResampler',
resampler_params=None,
bc=None,
bc_params=None,
):
## RUNNER
# choose which integrator to use
integrator_class = [
i for i in self.INTEGRATORS if i.__name__ == integrator
][0]
integrator_name = integrator_class.__name__
# use either the default params or the user params
if integrator_params is None:
integrator_params = self.DEFAULT_INTEGRATOR_PARAMS[integrator_name]
integrator = integrator_class(*integrator_params)
# TODO: not handling the params here
if platform_params is None:
platform_params = self.DEFAULT_PLATFORM_PARAMS[platform]
# handle additional runner options
if runner_params is None:
runner_params = {}
# make the runner for the test system
runner = OpenMMRunner(self.system,
self.topology,
integrator,
platform=platform,
**runner_params)
# RESAMPLER
# choose which resampler to use
resampler_class = [
res for res in self.RESAMPLERS if res.__name__ == resampler
][0]
resampler_name = resampler_class.__name__
# use either the default params or the user params
if resampler_params is None:
resampler_params = self.DEFAULT_RESAMPLER_PARAMS[resampler_name]
resampler = resampler_class(distance=self.distance,
init_state=self.init_state,
**resampler_params)
# BOUNDARY_CONDITIONS
# you don't have to specify a boundary condition
bc_name = bc
if bc_name is not None:
# choose which bc to use
bc_class = [res for res in self.BCS if res.__name__ == bc_name][0]
bc_name = bc_class.__name__
# use either the default params or the user params
if bc_params is None:
bc_params = self.DEFAULT_BC_PARAMS[bc_name]
bc = self.make_bc(bc_class, bc_params)
# APPARATUS
# build the apparatus
sim_apparatus = WepySimApparatus(runner,
resampler=resampler,
boundary_conditions=bc)
return sim_apparatus
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 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,
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)
# get the initial state from the context
MINIMIZED_INIT_OMM_STATE = simulation.context.getState(getPositions=True,
getVelocities=True,
getParameters=True,
getForces=True,
getEnergy=True,
getParameterDerivatives=True)
#### Wepy Orchestrator
### Apparatus
## 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
def test_lj_sim_manager_openmm_integration_run(
self,
class_tmp_path_factory,
boundary_condition_class,
resampler_class,
work_mapper_class,
platform,
lj_params,
lj_omm_sys,
lj_integrator,
lj_reporter_classes,
lj_reporter_kwargs,
lj_init_walkers,
lj_openmm_runner,
lj_unbinding_bc,
lj_wexplore_resampler,
lj_revo_resampler,
):
"""Run all combinations of components in the fixtures for the smallest
amount of time, just to make sure they all work together and don't give errors."""
logging.getLogger().setLevel(logging.DEBUG)
install_mp_handler()
logging.debug("Starting the test")
print("starting the test")
# the configuration class gives us a convenient way to
# parametrize our reporters for the locale
from wepy.orchestration.configuration import Configuration
# the runner
from wepy.runners.openmm import OpenMMRunner
# mappers
from wepy.work_mapper.mapper import Mapper
from wepy.work_mapper.worker import WorkerMapper
from wepy.work_mapper.task_mapper import TaskMapper
# the worker types for the WorkerMapper
from wepy.work_mapper.worker import Worker
from wepy.runners.openmm import OpenMMCPUWorker, OpenMMGPUWorker
# the walker task types for the TaskMapper
from wepy.work_mapper.task_mapper import WalkerTaskProcess
from wepy.runners.openmm import OpenMMCPUWalkerTaskProcess, OpenMMGPUWalkerTaskProcess
n_cycles = 1
n_steps = 2
num_workers = 2
# generate the reporters and temporary directory for this test
# combination
tmpdir_template = 'lj_fixture_{plat}-{wm}-{res}-{bc}'
tmpdir_name = tmpdir_template.format(plat=platform,
wm=work_mapper_class,
res=resampler_class,
bc=boundary_condition_class)
# make a temporary directory for this configuration to work with
tmpdir = str(class_tmp_path_factory.mktemp(tmpdir_name))
# make a config so that the reporters get parametrized properly
reporters = Configuration(
work_dir=tmpdir,
reporter_classes=lj_reporter_classes,
reporter_partial_kwargs=lj_reporter_kwargs).reporters
steps = [n_steps for _ in range(n_cycles)]
# choose the components based on the parametrization
boundary_condition = None
resampler = None
walker_fixtures = [lj_init_walkers]
runner_fixtures = [lj_openmm_runner]
boundary_condition_fixtures = [lj_unbinding_bc]
resampler_fixtures = [lj_wexplore_resampler, lj_revo_resampler]
walkers = lj_init_walkers
boundary_condition = [
boundary_condition
for boundary_condition in boundary_condition_fixtures
if type(boundary_condition).__name__ == boundary_condition_class
][0]
resampler = [
resampler for resampler in resampler_fixtures
if type(resampler).__name__ == resampler_class
][0]
assert boundary_condition is not None
assert resampler is not None
# generate the work mapper given the type and the platform
work_mapper_classes = {
mapper_class.__name__: mapper_class
for mapper_class in [Mapper, WorkerMapper, TaskMapper]
}
# # select the right one given the option
# work_mapper_type = [mapper_type for mapper_type in work_mapper_classes
# if type(mapper_type).__name__ == work_mapper_class][0]
# decide based on the platform and the work mapper which
# platform dependent components to build
if work_mapper_class == 'Mapper':
# then there is no settings
work_mapper = Mapper()
elif work_mapper_class == 'WorkerMapper':
if platform == 'CUDA' or platform == 'OpenCL':
work_mapper = WorkerMapper(num_workers=num_workers,
worker_type=OpenMMGPUWorker,
device_ids={
'0': 0,
'1': 1
},
proc_start_method='spawn')
if platform == 'OpenCL':
work_mapper = WorkerMapper(
num_workers=num_workers,
worker_type=OpenMMGPUWorker,
device_ids={
'0': 0,
'1': 1
},
)
elif platform == 'CPU':
work_mapper = WorkerMapper(
num_workers=num_workers,
worker_type=OpenMMCPUWorker,
worker_attributes={'num_threads': 1})
elif platform == 'Reference':
work_mapper = WorkerMapper(
num_workers=num_workers,
worker_type=Worker,
)
elif work_mapper_class == 'TaskMapper':
if platform == 'CUDA':
work_mapper = TaskMapper(
num_workers=num_workers,
walker_task_type=OpenMMGPUWalkerTaskProcess,
device_ids={
'0': 0,
'1': 1
},
proc_start_method='spawn')
elif platform == 'OpenCL':
work_mapper = TaskMapper(
num_workers=num_workers,
walker_task_type=OpenMMGPUWalkerTaskProcess,
device_ids={
'0': 0,
'1': 1
})
elif platform == 'CPU':
work_mapper = TaskMapper(
num_workers=num_workers,
walker_task_type=OpenMMCPUWalkerTaskProcess,
worker_attributes={'num_threads': 1})
elif platform == 'Reference':
work_mapper = TaskMapper(
num_workers=num_workers,
worker_type=WalkerTaskProcess,
)
else:
raise ValueError("Platform {} not recognized".format(platform))
# initialize the runner with the platform
runner = OpenMMRunner(lj_omm_sys.system,
lj_omm_sys.topology,
lj_integrator,
platform=platform)
logging.debug("Constructing the manager")
manager = Manager(walkers,
runner=runner,
boundary_conditions=boundary_condition,
resampler=resampler,
work_mapper=work_mapper,
reporters=reporters)
# since different work mappers need different process start
# methods for different platforms i.e. CUDA and linux fork
# vs. spawn we choose the appropriate one for each method.
logging.debug("Starting the simulation")
walkers, filters = manager.run_simulation(n_cycles,
steps,
num_workers=num_workers)
