def test_nearest_nbs(cut_off,frame_ind): nearest_nbs = [] for this_ind in test.water_inds[0:2]: nbs = md.compute_neighbors(test.traj[frame_ind], cut_off,[this_ind],haystack_indices = test.water_inds)[0] print this_ind while len(nbs)!=3: if len(nbs) > 3: pairs = [[this_ind,this_nb] for this_nb in nbs] distances=md.compute_distances(test.traj[frame_ind],pairs)[0] print distances min_three = f(distances,3) print min_three print nbs nbs = [nbs[ii] for ii in min_three] else: print 'increase cut_off!' nbs.append(this_ind) nbs.sort() if nbs in nearest_nbs: print "not unique" else: nearest_nbs.append(nbs) return np.array(nearest_nbs)
def main(): rec = args.rec r = mdtraj.load(rec) lig = args.lig l = mdtraj.load(lig) d = args.dist temp = args.temp pref = rec.split('.')[0] # can't figure out a non-hacky way to combine pdbs. cp_receptor = 'head -n -2 ' + rec receptor, r_err = call_cl(cp_receptor) cp_ligand = 'tail -n +2 ' + lig ligand, l_err = call_cl(cp_ligand) tf = open(temp, 'w') tf.write(receptor) tf.write(ligand) tf.close() # get indices of receptor within distance d of ligand comb = mdtraj.load(temp) # get ligand indices li = comb.topology.select('not protein') # find neighbors neighbors = mdtraj.compute_neighbors(comb, d, li)[0] # remove ligand from neighbor list (symmetric diff) n = np.setxor1d(li, neighbors) # easier to reset sometimes comb = mdtraj.load(temp) comb.restrict_atoms(n) comb.save_pdb(pref + '_trim.pdb') IPython.embed()
def test_compute_neighbors_3(): traj = md.load(get_fn('test_good.nc'), top=get_fn('test.parm7')) query_indices = traj.top.select('residue 1') cutoff = 1.0 value = md.compute_neighbors(traj, cutoff, query_indices) reference = compute_neighbors_reference(traj, cutoff, query_indices) for i in range(traj.n_frames): eq(value[i], reference[i])
def test_compute_neighbors_2(get_fn): traj = md.load(get_fn('4ZUO.pdb')) query_indices = traj.top.select('residue 1') cutoff = 1.0 value = md.compute_neighbors(traj, cutoff, query_indices) reference = compute_neighbors_reference(traj, cutoff, query_indices) for i in range(traj.n_frames): eq(value[i], reference[i])
def find_nearest_nbs(self,cut_off,frame_ind,N_nbs): """Finds the N nearest neighbors of all waters in a single simulation frame and returns all unique nearest-neighbor tthds as an (n,4) array. n is the number of tthds found in this frame. Every tthd is represented by 4 indices of the water molecules (oxygens) that are the for vertices. More generally, this function be used to form ensembles with (N_nbs+1) vertices. The numpy array returned will have shape (n, N_nbs+1). Parameters ---------- cut_off : float cuttoff distance for looking for nearest neighbors. Should be set such that at least N_nbs waters can be found within that distance frame_ind : int the index of the frame in which to find nearest neighbors N_nbs : int the number of nearest neighbors to return. set to 3 in this case since we are interest in 4-point correlators. """ nearest_nbs = [] frame = self.traj[frame_ind] for this_ind in self.water_inds: # find all the neighbors with the cut_off distance for each water molecule nbs = md.compute_neighbors(frame, cut_off,[this_ind],haystack_indices = self.water_inds)[0] while len(nbs)!= N_nbs: if len(nbs) > N_nbs: # compute all the distance between neighbors and only keep the closest # N_nbs ones pairs = [[this_ind,this_nb] for this_nb in nbs] distances=md.compute_distances(frame,pairs)[0] min_three = np.argsort(distances)[:][:N_nbs] nbs = [nbs[ii] for ii in min_three] else: print 'increase cut_off!' # keep nbs a list so append can happen nbs = [nbs[ii] for ii in range(len(nbs))] # add the ind of the water for whom neighbors have been found to the list nbs.append(this_ind) # sort the list so it can be comapred to set of neighbors we already found nbs.sort() if nbs in nearest_nbs: # if this set of nearest neighbors already exist, don't added it to the list print "not unique" else: # only add to list if unique nearest_nbs.append(nbs) # return list as an numpy array return np.array(nearest_nbs)
def test_compute_neighbors_1(): n_frames = 2 n_atoms = 20 cutoff = 2 xyz = random.randn(n_frames, n_atoms, 3) traj = md.Trajectory(xyz=xyz, topology=None) query_indices = [0, 1] value = md.compute_neighbors(traj, cutoff, query_indices) reference = compute_neighbors_reference(traj, cutoff, query_indices) for i in range(n_frames): eq(value[i], reference[i])
def _neighbouring_atoms(md_topology, trajectory, atom_subset, atom_number, verbose, unpythonize, chunk, cutoff): # first create a list with all the paths that are needed try: trajectory_path = os.listdir(trajectory) except: sys.exit('Make sure you have provided a string for a valid path to a trajectory file!') else: if verbose > 0: print 'Loading trajectories from the following files: ' for trajectory_i in trajectory_path: print trajectory_i # initiate some variables neighbour_atoms = [] sim_time=[] number_of_frames = 0 # now we need to load each trajectory file as a chunk try: pbar = tqdm(total=len(trajectory_path), unit= 'File') for file_i in trajectory_path: for chunk_i in md.iterload(trajectory + file_i, chunk, top=md_topology, atom_indices = atom_subset): sim_time.append(chunk_i.time) number_of_frames += chunk_i.n_frames if verbose > 1: print 'Successfully loaded trajectory: \n %s' %(chunk_i) neighbour_atoms.append(md.compute_neighbors(chunk_i, cutoff, np.array([atom_number]))) neighbour_atoms_np =np.concatenate(neighbour_atoms) pbar.update(1) except: sys.exit('Make sure you provided a valid path to a folder with trajectory files!') else: print '\nSuccesfully loaded coordinates for %s atoms in %s frames!' %(len(atom_subset), number_of_frames) all_neighbour_atoms_np = np.concatenate(neighbour_atoms_np) sim_time = np.concatenate(sim_time) return all_neighbour_atoms_np, sim_time
def make_tthd(self,vertex_ind,cut_off,frame_ind): """ Given index of one vertex, find three other vertices within radius cut_off to form a unique tetrahedron in one single frame of traj returns list a array of shape (3,), indices of the three other vertices """ nbs = md.compute_neighbors(self.traj[frame_ind],cut_off,[vertex_ind],haystack_indices = self.water_inds)[0] # print nbs.shape tthd_inds = np.array(list(combinations(nbs,3))) # I should not need to sort nbs, consider getting rid of the sorted part for all other instances of combinations tthds = [] xyz_pos = self.traj[frame_ind].xyz for this_tthd in tthd_inds: # representing a tetrahedron with just two vectors, is this even right? r_ij = xyz_pos[0,vertex_ind,:]- xyz_pos[0,this_tthd[0],:] r_kl = xyz_pos[0,this_tthd[1],:]- xyz_pos[0,this_tthd[2],:] # nm tthds.append([r_ij,r_kl]) return tthds
def visualize_complex(complex_mdtraj): ligand_atoms = [a.index for a in complex_mdtraj.topology.atoms if "LIG" in str(a.residue)] binding_pocket_atoms = md.compute_neighbors(complex_mdtraj, 0.5, ligand_atoms)[0] binding_pocket_residues = list(set([complex_mdtraj.topology.atom(a).residue.resSeq for a in binding_pocket_atoms])) binding_pocket_residues = [str(r) for r in binding_pocket_residues] binding_pocket_residues = " or ".join(binding_pocket_residues) traj = nglview.MDTrajTrajectory( complex_mdtraj ) # load file from RCSB PDB ngltraj = nglview.NGLWidget( traj ) ngltraj.representations = [ { "type": "cartoon", "params": { "sele": "protein", "color": "residueindex" } }, { "type": "licorice", "params": { "sele": "(not hydrogen) and (%s)" % binding_pocket_residues } }, { "type": "ball+stick", "params": { "sele": "LIG" } } ] return ngltraj
def make_pairs(self,vertex_ind,cut_off,frame_ind): """Returns a list of vectors that connect pairs of water molecules that are within a cutoff distances of one water molecules Parameters ---------- vertex_ind : int index of the water molecule whose neighbors one want to find and form pairs with cut_off : float distance in nm within which to search for neighbors frame_ind : int index of simulation in which to make pairs """ nbs = md.compute_neighbors(self.traj[frame_ind],cut_off,[vertex_ind],haystack_indices = self.water_inds)[0] pairs = [] xyz_pos = self.traj[frame_ind].xyz for this_nb in nbs: r_ij = xyz_pos[0,vertex_ind,:]- xyz_pos[0,this_nb,:] pairs.append(r_ij) return pairs
def nb_finder(cut_off,vertex_ind,frame_ind): nbs = md.compute_neighbors(test.traj[frame_ind], cut_off,[vertex_ind],haystack_indices = test.water_inds)[0] return nbs
def test_compute_neighbors_2(): traj = md.load(get_fn('4ZUO.pdb')) query_indices = traj.top.select('residue 1') cutoff = 1.0 value = md.compute_neighbors(traj, cutoff, query_indices) reference = compute_neighbors_reference(traj, cutoff, query_indices)