-
-
Save kyleabeauchamp/6136191 to your computer and use it in GitHub Desktop.
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) |
PS Osama, some of those distances are going to be constant during the MD simulation, as we typically constrain the bond lengths. If you like, you can try to make a better script that deletes those atompairs.
If speed is an issue and you have a lot of frames of MD, you can also use msmbuilder.geometry.contact.atom_distances
, which is the same function as mdtraj.geometry.distance.compute_distance
, but probably a little faster. I haven't ported that C code to mdtraj yet.
And this requires mdtraj version 0.4.0, which was just released yesterday. you should be able to upgrade with pip install mdtraj --upgrade
if you haven't already.
PS: Robert it's super convenient having access to the element column, which I never had in any previous PDB readers...
I added an optional conversion of the distance matrix into DataFrame...
Here is the preliminary code to load a trajectory, find the desired pairs of atom indices, and calculate the distances.