python - 将残留物 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