pythonnumpymoleculemdanalysis

Adding residue IDs to a numpy array consisting of time series data of water coordinates


I got this script for generating time series data of water molecules, and I want to add one more header row to that generated matrix with residue IDs of water molecules. Could anybody help with with reworking this script? Thanks!

import numpy as np
import MDAnalysis as mda

u = mda.Universe(PSF, DCD)
water_oxygens = u.select_atoms("name OW")

# pre-allocate the array for the data
data = np.zeros((u.trajectory.n_frames, water_oxygens.n_atoms + 1))

for i, ts in enumerate(u.trajectory):
   data[i, 0] = ts.time                          # store current time
   data[i, 1:] = water_oxygens.positions[:, 2]   # extract all z-coordinates

Solution

  • Here is an adjusted code example. You might need to install the package MDAnalysisTests to run it:

    import numpy as np
    import MDAnalysis as mda
    from MDAnalysisTests.datafiles import waterPSF, waterDCD
    
    u = mda.Universe(waterPSF, waterDCD)
    water_oxygens = u.select_atoms("name OH2")
    
    # pre-allocate the array for the data
    # one extra row for the header water residue IDs
    data = np.zeros((u.trajectory.n_frames + 1, water_oxygens.n_atoms + 1))
    
    # initialise the water residue IDs
    data[0, 0] = np.NaN # the time column
    data[0, 1:] = water_oxygens.atoms.resids
    
    for i, ts in enumerate(u.trajectory, start=1):
       data[i, 0] = ts.time                          # store current time
       data[i, 1:] = water_oxygens.positions[:, 2]   # extract all z-coordinates