将剩余ID添加到由水坐标的时间序列数据组成的numpy数组中

Question

我得到了这个脚本，用于生成水分子的时间序列数据，我想在生成的矩阵中再添加一个标题行，其中包含水分子的剩余ID。有人能帮忙修改这个脚本吗？谢谢

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

下面是一个调整后的代码示例。您可能需要安装程序包MDAnalysisTests才能运行它：

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