Pandas Trajectory Container

pandas_traj.py

import mdtraj, mdtraj.geometry
import pandas as pd
import numpy as np

r = mdtraj.load(filename)

data = []
for atom in r.top.atoms:
    data.append((atom.index, atom.name, atom.residue.index, atom.residue.name, atom.residue.chain.index))

T = pd.DataFrame(data, columns = ["index", "atom", "rid" , "residue", "chain"])

T[T.residue == "CYS"][0:10]

In [9]: T[T.residue == "CYS"][0:10]
Out[9]: 
    index atom  rid residue  chain
9       9    N    1     CYS      0
10     10    H    1     CYS      0
11     11   CA    1     CYS      0
12     12   HA    1     CYS      0
13     13   CB    1     CYS      0
14     14  HB3    1     CYS      0
15     15  HB2    1     CYS      0
16     16   SG    1     CYS      0
17     17    C    1     CYS      0
18     18    O    1     CYS      0

replacements = {"CYS":"CYSS"}

def map_residue(name):
    if name not in replacements.keys():
        return name
    else:
        return replacements[name]

T.residue = T.residue.apply(map_residue)

T[T.residue == "CYSS"][0:10]

In [13]: T[T.residue == "CYSS"][0:10]
Out[13]: 
    index atom  rid residue  chain
9       9    N    1    CYSS      0
10     10    H    1    CYSS      0
11     11   CA    1    CYSS      0
12     12   HA    1    CYSS      0
13     13   CB    1    CYSS      0
14     14  HB3    1    CYSS      0
15     15  HB2    1    CYSS      0
16     16   SG    1    CYSS      0
17     17    C    1    CYSS      0
18     18    O    1    CYSS      0