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)