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)
