def init(self, **kwargs): """ Parameters ---------- **kwargs : Returns ------- """ super().init(**kwargs) # load the json topology as an mdtraj one mdtraj_top = json_to_mdtraj_topology(self.json_main_rep_top) # make a traj for the initial state to use as a topology for # visualizing the walkers init_traj = mdj.Trajectory( [self.init_main_rep_positions], unitcell_lengths=[self.init_unitcell_lengths], unitcell_angles=[self.init_unitcell_angles], topology=mdtraj_top) # write out the init traj as a pdb logging.info("Writing initial state to {}".format( self.init_state_path)) init_traj.save_pdb(self.init_state_path)
def binding_site_idxs(json_topology, coords, box_vectors, cutoff): # convert quantities to numbers in nanometers cutoff = cutoff.value_in_unit(unit.nanometer) coords = coords.value_in_unit(unit.nanometer) box_lengths, box_angles = box_vectors_to_lengths_angles( box_vectors.value_in_unit(unit.nanometer)) # selecting ligand and protein binding site atom indices for # resampler and boundary conditions lig_idxs = ligand_idxs(json_topology) prot_idxs = protein_idxs(json_topology) # make a trajectory to compute the neighbors from traj = mdj.Trajectory(np.array([coords]), unitcell_lengths=[box_lengths], unitcell_angles=[box_angles], topology=json_to_mdtraj_topology(json_topology)) # selects protein atoms which have less than 8 A from ligand # atoms in the crystal structure neighbors_idxs = mdj.compute_neighbors(traj, cutoff, lig_idxs) # selects protein atoms from neighbors list binding_selection_idxs = np.intersect1d(neighbors_idxs, prot_idxs) return binding_selection_idxs
def _calc_min_distance(self, walker): """Min-min distance for a walker. Parameters ---------- walker : object implementing the Walker interface Returns ------- min_distance : float """ cell_lengths, cell_angles = box_vectors_to_lengths_angles( walker.state['box_vectors']) # convert the json topology to an mdtraj one mdj_top = json_to_mdtraj_topology(self._topology) # make a traj out of it so we can calculate distances through # the periodic boundary conditions walker_traj = mdj.Trajectory(walker.state['positions'], topology=mdj_top, unitcell_lengths=cell_lengths, unitcell_angles=cell_angles) # calculate the distances through periodic boundary conditions # and get hte minimum distance min_distance = np.min( mdj.compute_distances( walker_traj, it.product(self.ligand_idxs, self.receptor_idxs))) return min_distance
def init(self, **kwargs): """ Parameters ---------- **kwargs : Returns ------- """ super().init(**kwargs) if self.init_image is not None: image_mdj_topology = json_to_mdtraj_topology( self.json_main_rep_top) # initialize the initial image into the image traj init_image_traj = mdj.Trajectory([self.init_image], time=self.times, topology=image_mdj_topology) # save this as a PDB for a topology to view in VMD etc. to go # along with the trajectory we will make logging.info("Writing initial image to {}".format( self.init_state_path)) init_image_traj.save_pdb(self.init_state_path) self._top_pdb_written = True
def report(self, cycle_idx=None, new_walkers=None, **kwargs): """Report the current cycle's walker states as 3D molecular structures. Parameters ---------- cycle_idx : int new_walkers : list of Walker objects """ # load the json topology as an mdtraj one mdtraj_top = json_to_mdtraj_topology(self.json_main_rep_top) # slice off the main_rep indices because that is all we want # to write for these main_rep_positions = np.array([ walker.state['positions'][self.main_rep_idxs] for walker in new_walkers ]) # convert the box vectors unitcell_lengths, unitcell_angles = traj_box_vectors_to_lengths_angles( np.array([walker.state['box_vectors'] for walker in new_walkers])) # make a trajectory from these walkers traj = mdj.Trajectory(main_rep_positions, unitcell_lengths=unitcell_lengths, unitcell_angles=unitcell_angles, topology=mdtraj_top) # write to the file for this trajectory traj.save_dcd(self.walker_path)
def report(self, cycle_idx=None, resampler_data=None, **kwargs): # load the json topology as an mdtraj one image_mdj_topology = json_to_mdtraj_topology(self.json_main_rep_top) # collect the new images defined new_images = [] for resampler_rec in resampler_data: image = resampler_rec['image'] new_images.append(image) times = np.array([cycle_idx + 1 for _ in range(len(new_images))]) # combine the new image positions and times with the old self.image_traj_positions.extend(new_images) self.times.extend(times) # only save if we have an image yet if len(self.image_traj_positions) > 0: # make a trajectory of the new images, using the cycle_idx as the time new_image_traj = mdj.Trajectory(self.image_traj_positions, time=self.times, topology=image_mdj_topology) # if we haven't already written a topology PDB write it now if not self._top_pdb_written: new_image_traj[0].save_pdb(self.init_state_path) self._top_pdb_written = True # then the images to the trajectory file new_image_traj.save_dcd(self.image_path)
def binding_site_idxs(json_topology, ligand_idxs, receptor_idxs, coords, box_vectors, cutoff, periodic=True): """ Parameters ---------- json_topology : str ligand_idxs : arraylike (1,) receptor_idxs : arraylike (1,) coords : simtk.Quantity box_vectors : simtk.Quantity cutoff : float Returns ------- binding_site_idxs : arraylike (1,) """ # convert quantities to numbers in nanometers cutoff = cutoff.value_in_unit(unit.nanometer) coords = coords.value_in_unit(unit.nanometer) box_lengths, box_angles = box_vectors_to_lengths_angles( box_vectors.value_in_unit(unit.nanometer)) # make a trajectory to compute the neighbors from traj = mdj.Trajectory(np.array([coords]), unitcell_lengths=[box_lengths], unitcell_angles=[box_angles], topology=json_to_mdtraj_topology(json_topology)) neighbors_idxs = mdj.compute_neighbors(traj, cutoff, ligand_idxs, periodic=periodic)[0] # selects protein atoms from neighbors list binding_selection_idxs = np.intersect1d(neighbors_idxs, receptor_idxs) return binding_selection_idxs
def init(self, **kwargs): """Initialize the reporter at simulation time. This will generate the initial state PDB file. """ super().init(**kwargs) # load the json topology as an mdtraj one mdtraj_top = json_to_mdtraj_topology(self.json_main_rep_top) # make a traj for the initial state to use as a topology for # visualizing the walkers init_traj = mdj.Trajectory( [self.init_main_rep_positions], unitcell_lengths=[self.init_unitcell_lengths], unitcell_angles=[self.init_unitcell_angles], topology=mdtraj_top) # write out the init traj as a pdb logging.info("Writing initial state to {}".format( self.init_state_path)) init_traj.save_pdb(self.init_state_path)
def binding_site_idxs(json_topology, ligand_idxs, receptor_idxs, coords, box_vectors, cutoff, periodic=True): """Parameters ---------- json_topology : str ligand_idxs : arraylike (1,) receptor_idxs : arraylike (1,) coords : N x 3 arraylike of float or simtk.Quantity If not a quantity will implicitly be treated as being in nanometers. box_vectors : simtk.Quantity If not a quantity will implicitly be treated as being in nanometers. cutoff : float or simtk.Quantity If not a quantity will implicitly be treated as being in nanometers. Returns ------- binding_site_idxs : arraylike (1,) """ # if they are simtk.units convert quantities to numbers in # nanometers if unit.is_quantity(cutoff): cutoff = cutoff.value_in_unit(unit.nanometer) if unit.is_quantity(coords): coords = coords.value_in_unit(unit.nanometer) if unit.is_quantity(box_vectors): box_vectors = box_vectors.value_in_unit(unit.nanometer) box_lengths, box_angles = box_vectors_to_lengths_angles(box_vectors) # make a trajectory to compute the neighbors from traj = mdj.Trajectory(np.array([coords]), unitcell_lengths=[box_lengths], unitcell_angles=[box_angles], topology=json_to_mdtraj_topology(json_topology)) neighbors_idxs = mdj.compute_neighbors(traj, cutoff, ligand_idxs, periodic=periodic)[0] # selects protein atoms from neighbors list binding_selection_idxs = np.intersect1d(neighbors_idxs, receptor_idxs) return binding_selection_idxs
def __init__(self, initial_state=None, cutoff_distance=1.0, topology=None, ligand_idxs=None, receptor_idxs=None, periodic=True, **kwargs): """Constructor for UnbindingBC class. All the key-word arguments are necessary. The 'initial_state' should be the initial state of your simulation for proper non-equilibrium simulations. Arguments --------- initial_state : object implementing State interface The state walkers will take on after unbinding. cutoff_distance : float The distance that specifies the boundary condition. When the min-min ligand-receptor distance is less than this it will be warped. topology : str A JSON string of topology. ligand_idxs : list of int Indices of the atoms in the topology that correspond to the ligands. receptor_idxs : list of int Indices of the atoms in the topology that correspond to the receptor for the ligand. Raises ------ AssertionError If any of the following are not provided: initial_state, topology, ligand_idxs, receptor_idxs AssertionError If the cutoff distance is not a float. Warnings -------- The 'initial_state' should be the initial state of your simulation for proper non-equilibrium simulations. Notes ----- The topology argument is necessary due to an implementation detail that uses mdtraj and may not be required in the future. """ # since the super class can handle multiple initial states we # wrap the single initial state to a list. super().__init__(initial_states=[initial_state], ligand_idxs=ligand_idxs, receptor_idxs=receptor_idxs, **kwargs) # test input assert topology is not None, "Must give a reference topology" assert type(cutoff_distance) is float self._cutoff_distance = cutoff_distance self._topology = topology # convert the json topology to an mdtraj one self._mdj_top = json_to_mdtraj_topology(self._topology) # whether or not to use the periodic box vectors in the # distance calculation self._periodic = periodic
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_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)