Skip to content

Instantly share code, notes, and snippets.

@kyleabeauchamp

kyleabeauchamp/atom_pair_distances.py

Last active July 20, 2021 13:10
Show Gist options
  • Star 1 You must be signed in to star a gist
  • Fork 1 You must be signed in to fork a gist
  • Save kyleabeauchamp/6136191 to your computer and use it in GitHub Desktop.
Save kyleabeauchamp/6136191 to your computer and use it in GitHub Desktop.
Download ZIP
Calculate specific atom pair distances with MDTraj
Raw
atom_pair_distances.py
import itertools
import mdtraj, mdtraj.geometry
import pandas as pd
traj = mdtraj.load("./traj.xtc", top="native.pdb")
top, bonds = traj.top.to_dataframe()
atoms = np.array(["H", "HA", "N", "CA", "C", "CB"])
bad_residues = np.array(["GLY", "PRO"])
top = top[np.in1d(top.name, atoms)] # Selected desired atoms
top = top[~np.in1d(top.resName, bad_residues)] # Remove unwanted residues
atom_pairs = []
for k, x in top.groupby("resSeq"):
atom_pairs.extend(list(itertools.combinations(x.index, 2)))
atom_pairs = np.array(atom_pairs)
distances = mdtraj.geometry.distance.compute_distances(traj, atom_pairs)
id0 = top.ix[atom_pairs[:,0]].resSeq
id1 = top.ix[atom_pairs[:,1]].resSeq
name0 = top.ix[atom_pairs[:,0]].name
name1 = top.ix[atom_pairs[:,1]].name
columns = pd.MultiIndex.from_arrays([id0, name0, id1, name1])
distances = pd.DataFrame(distances, columns=columns)
@kyleabeauchamp
Copy link
Author

PS: Robert it's super convenient having access to the element column, which I never had in any previous PDB readers...

@kyleabeauchamp
Copy link
Author

I added an optional conversion of the distance matrix into DataFrame...

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment