Search code examples
pythonmdanalysis

How to iterate efficiently through MDAnalysis trajectory and save residue attribute time series?


I have some working code using MDAnalysis that saves the residues center of mass time series in an array, but I wonder if there is a more Pythonic or overall efficient/fast way (comprehensions, array operations...) to do it.

import MDAnalysis as mda
mdau = mda.Universe(pdb, xtc)
arr = np.empty(( len(mdau.select_atoms("protein").residues), len(mdau.trajectory), 3 ))
# 288 protein residues, 1250 frames and 3 xyz-coordinates per center of mass; this array shape is important
for ts in mdau.trajectory:
    for num, res in enumerate(mdau.select_atoms("protein").residues):
        arr[num, ts.frame] = res.atoms.center_of_mass()

The .pdb and .xtc files I am using can be downloaded in these links: https://submission.gpcrmd.org/dynadb/files/Dynamics/11579_dyn_169.pdb https://submission.gpcrmd.org/dynadb/files/Dynamics/11576_trj_169.xtc


Solution

  • There are a couple of changes you could make to your code:

    • select your protein atoms outside of the for loop, rather than on each iteration
    • vectorise the center of mass calculation over residues by passing the compound='residues' argument to the center_of_mass method
    • use the ag.n_residues and u.trajectory.n_frames attributes

    Here's an update to your code that uses these suggestions:

    import numpy as np
    import MDAnalysis as mda
    
    u = mda.Universe('11579_dyn_169.pdb', '11576_trj_169.xtc')
    protein = u.select_atoms("protein")
    arr = np.empty((protein.n_residues, u.trajectory.n_frames, 3))
    
    for ts in u.trajectory:
        arr[:, ts.frame] = protein.center_of_mass(compound='residues')