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

Question

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

Answer 1

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