<think>这个错误是因为在创建cKDTree时,传入的坐标数据有负值,而设置了boxsize(周期性边界条件)时,要求坐标必须在[0, boxsize]范围内。
根据错误信息,我们的坐标数据有负值,而我们的盒子是从0到100,所以负值超出了盒子范围。
解决方案:
1. 在创建cKDTree之前,将坐标映射到盒子范围内(即应用周期性边界条件,将负坐标加上盒子长度,直到在[0, boxsize]内)。
2. 或者,在读取文件后,对坐标进行周期修正。
但是,注意在LAMMPS中,坐标是可以超出盒子范围的,但在进行空间搜索时需要应用周期性边界条件。cKDTree的boxsize参数要求坐标在[0, boxsize]内。
因此,我们需要在构建树之前将坐标映射到盒子内。
修改方案:
在函数`identify_atom_types`中构建KDTree之前,对坐标进行周期性修正。
具体步骤:
1. 复制坐标数组,避免修改原始数据。
2. 对每个维度,应用周期性边界条件:将坐标减去盒子长度*floor(坐标/盒子长度),即映射到[0, boxsize)区间。
但是,注意我们的盒子尺寸box是一个长度为3的数组,表示三个维度的盒子长度。
我们可以写一个函数来将坐标映射到盒子内:
另外,在计算角度和距离时,我们已经在使用分数坐标来处理周期性,但是在构建树时,树需要坐标在[0, boxsize]内。
因此,我们在`identify_atom_types`函数中,构建树之前,先修正坐标到盒子内。
注意:这个修正只用于构建树,不改变原始坐标(因为原始坐标在后续处理中也要注意周期性)。
但是,在`identify_atom_types`函数中,我们传入的coords是原始坐标,所以我们需要在函数内部进行修正,并用于构建树。
由于在构建树后,我们使用树进行邻居搜索,而邻居搜索会考虑周期性(通过boxsize参数),但要求坐标在[0, boxsize]内。所以我们必须先修正坐标。
因此,我们在`identify_atom_types`函数中增加一步:将坐标映射到盒子内。
修正坐标的函数可以这样写:
但是注意:在`read_data_file`和`read_lammpstrj_frame`中,我们读取的坐标可能是任意值,所以我们需要在构建树之前进行修正。
我们可以在`identify_atom_types`函数内部进行修正,因为只有这个函数需要构建树。
修正代码:
在构建树之前:
coords = coords % box # 这样不行,因为box是三维向量,且每个维度不同
正确做法:
for i in range(3):
coords[:, i] = coords[:, i] - box[i] * np.floor(coords[:, i] / box[i])
但是,注意:这样会改变传入的coords,而原始坐标在后续处理中(比如角度计算)也需要在同一个周期内?为了避免混淆,我们在函数内使用修正后的坐标来构建树,但原始坐标保持不变。
但是,在邻居搜索中,我们使用的是修正后的坐标,而后续计算角度时使用的是原始坐标(未修正)。这样可能导致问题,因为两个点的坐标可能不在同一个周期内。
因此,我们需要在构建树之前修正坐标,并且在计算角度和距离时,也要注意使用相同的周期性处理。但是,在计算角度和距离时,我们已经使用了分数坐标和周期性修正(通过分数坐标差再映射到[-0.5,0.5]区间)。所以,在构建树时使用修正后的坐标,而在其他计算中仍然使用原始坐标(但通过分数坐标处理周期性)是可以的。
但是,为了确保一致性,我们最好在构建树时使用修正后的坐标,而在其他所有地方都使用原始坐标(因为我们的角度计算函数已经考虑了周期性)。因为邻居搜索返回的索引是原始坐标的索引,我们通过索引取原始坐标进行计算。
所以,我们在`identify_atom_types`函数中构建树之前,先对坐标进行周期性修正,然后使用修正后的坐标构建树。这样,邻居搜索会正确进行。
修改`identify_atom_types`函数:
在构建树之前添加:
# 修正坐标到盒子内 [0, box[i]) 区间
corrected_coords = coords.copy()
for dim in range(3):
corrected_coords[:, dim] = corrected_coords[:, dim] - box[dim] * np.floor(corrected_coords[:, dim] / box[dim])
# 然后使用corrected_coords构建树
注意:box是三维数组,表示每个维度的长度。
另外,在后续的邻居搜索中,我们使用corrected_coords来构建树,但是在计算距离时,我们使用原始坐标(因为邻居搜索返回的是索引,我们用索引去原始坐标数组中取点计算距离)。这样会有问题吗?
实际上,在构建树时,我们使用修正后的坐标,这样树的结构是正确的。然后,当我们查询邻居时,查询点也要用修正后的坐标。但是,在函数内,我们使用原始坐标来计算距离(比如在找H的供体时,我们使用原始坐标计算距离)。这样会导致距离计算错误,因为两个点可能不在同一个周期内。
因此,我们需要在计算距离时也使用修正后的坐标?或者,在计算距离时,我们使用原始坐标,但通过最小镜像约定(即计算周期性距离)来得到最短距离。
但是,cKDTree在构建时指定了boxsize,那么查询时也会考虑周期性,返回的距离是周期性最短距离。所以,我们构建树时使用修正后的坐标,并指定boxsize,这样查询时就会自动考虑周期性。
所以,在`identify_atom_types`函数中,我们使用修正后的坐标构建树,然后使用这个树进行邻居搜索(查询时使用修正后的坐标)。这样,返回的邻居距离是周期性的最短距离。
但是,在后续的供体选择中,我们使用原始坐标计算距离,这就不一致了。
因此,我们需要统一:在`identify_atom_types`函数中,所有使用树进行邻居搜索的地方,都要用修正后的坐标。但是,在计算距离选择最近供体时,我们使用原始坐标计算,这就不对了。
所以,我们有两个选择:
1. 在`identify_atom_types`函数中,使用修正后的坐标进行所有计算(包括距离计算)。
2. 或者,在构建树时,我们不修正坐标,而是要求坐标已经在[0, box]内(但我们现在有负值,所以不行)。
因此,我们选择在`identify_atom_types`函数内部,使用修正后的坐标进行所有计算(包括构建树和距离计算)。但是,这样会改变函数内的局部坐标,而原始坐标不变。但是,在函数内,我们只使用坐标来构建树和计算距离,这些计算应该使用修正后的坐标。
具体修改:
在`identify_atom_types`函数中:
# 修正坐标
corrected_coords = coords.copy()
for dim in range(3):
corrected_coords[:, dim] = corrected_coords[:, dim] - box[dim] * np.floor(corrected_coords[:, dim] / box[dim])
# 然后,在函数内,所有使用坐标的地方,都使用corrected_coords
但是,注意在函数中,我们使用坐标的地方有:
- 构建树:使用corrected_coords
- 在邻居搜索后,计算距离:使用corrected_coords(因为树是用corrected_coords构建的,查询点也要用corrected_coords,计算距离也要用corrected_coords)
但是,在计算距离时,我们使用np.linalg.norm(corrected_coords[a]-corrected_coords[b]),这样计算的是在修正后坐标空间的距离,由于修正后的坐标在同一个盒子内,所以这个距离就是最短距离(不需要再考虑周期性,因为已经修正到同一个周期内)。
因此,我们在函数内部使用corrected_coords代替原始coords。
但是,注意:在函数外部,原始坐标保持不变。在函数返回后,其他函数(如`find_hbonds_in_frame`)仍然使用原始坐标。这样,在角度计算中,我们使用原始坐标,但通过分数坐标处理周期性,所以应该没问题。
所以,修改`identify_atom_types`函数:
步骤:
1. 复制一份坐标,进行周期性修正。
2. 在函数内,使用修正后的坐标进行所有与距离相关的计算(包括构建树和邻居搜索)。
注意:在识别磷酸基团时,我们使用P原子周围的O原子,这个距离计算使用修正后的坐标。
修改后的`identify_atom_types`函数代码:
另外,在`find_hbonds_in_frame`函数中,我们同样需要构建树,所以同样需要修正坐标。因此,我们最好在`find_hbonds_in_frame`函数中也进行同样的坐标修正。
但是,为了避免重复代码,我们可以写一个辅助函数来修正坐标。
但是,由于两个函数都需要,我们写一个工具函数:
def wrap_coordinates(coords, box):
wrapped_coords = coords.copy()
for dim in range(3):
wrapped_coords[:, dim] = wrapped_coords[:, dim] - box[dim] * np.floor(wrapped_coords[:, dim] / box[dim])
return wrapped_coords
然后,在`identify_atom_types`和`find_hbonds_in_frame`函数中,都使用这个函数来修正坐标。
但是,注意在`find_hbonds_in_frame`函数中,我们构建了两次树:一次是全局树,一次是受体树。所以,在构建树之前,我们需要修正坐标。
因此,我们修改两个函数:
- 在`identify_atom_types`中,构建树之前,修正坐标。
- 在`find_hbonds_in_frame`中,构建树之前,修正坐标。
注意:在`find_hbonds_in_frame`函数中,我们使用全局树和受体树,都需要修正坐标。
但是,在计算角度时,我们使用的是原始坐标(因为角度计算函数已经考虑了周期性,使用分数坐标和晶格矩阵)。所以,在`find_hbonds_in_frame`函数中,我们使用原始坐标计算角度,而使用修正后的坐标进行邻居搜索。
这样,在`find_hbonds_in_frame`中,我们使用修正后的坐标构建树,然后进行邻居搜索,但计算角度时使用原始坐标。这样是可行的,因为角度计算函数内部已经处理了周期性。
所以,我们添加一个辅助函数:
```python
def wrap_coordinates(coords, box):
"""将坐标映射到周期性盒子内 [0, box] 区间"""
wrapped_coords = coords.copy()
for dim in range(3):
# 避免除以零
if box[dim] <= 0:
continue
wrapped_coords[:, dim] = wrapped_coords[:, dim] - box[dim] * np.floor(wrapped_coords[:, dim] / box[dim])
return wrapped_coords
```
然后,在`identify_atom_types`和`find_hbonds_in_frame`函数中,构建树之前,调用这个函数修正坐标。
但是,注意:在`identify_atom_types`函数中,我们使用`coords`(原始坐标)来构建树,现在改为使用修正后的坐标。在函数内部,所有使用坐标的地方(比如计算距离)都使用修正后的坐标。
同样,在`find_hbonds_in_frame`函数中,构建全局树和受体树时,使用修正后的坐标。
但是,在`find_hbonds_in_frame`函数中,我们使用原始坐标计算角度,这是正确的,因为角度计算函数`calculate_angle`内部使用分数坐标和晶格矩阵,已经考虑了周期性。
因此,我们修改这两个函数:
在`identify_atom_types`函数中:
corrected_coords = wrap_coordinates(coords, box)
kdtree = cKDTree(corrected_coords, boxsize=box) # 注意:boxsize需要是三个维度的盒子尺寸
然后,在查询邻居时,使用corrected_coords。
但是,在计算距离时,比如在找H原子的供体时,我们使用原始坐标计算距离?不,在`identify_atom_types`函数中,我们使用修正后的坐标计算距离,因为邻居搜索返回的索引对应的坐标是修正后的坐标,所以计算距离应该用修正后的坐标。
所以,在`identify_atom_types`函数中,计算距离的部分,我们使用corrected_coords。
修改`identify_atom_types`函数中计算距离的部分:
将:
dist = np.linalg.norm(coords[h_idx] - coords[o_idx])
改为:
dist = np.linalg.norm(corrected_coords[h_idx] - corrected_coords[o_idx])
同样,在`find_hbonds_in_frame`函数中,计算距离时也要使用修正后的坐标。
所以,我们在`find_hbonds_in_frame`函数中也做类似修改。
但是,注意:在`find_hbonds_in_frame`函数中,我们使用全局树(用修正后的坐标构建)进行查询,查询点也要用修正后的坐标。所以,在查询H原子的供体时,我们使用H原子的修正后坐标。
因此,在`find_hbonds_in_frame`函数中,我们也要修正坐标,并在计算距离时使用修正后的坐标。
修改步骤:
1. 在`identify_atom_types`函数开头,添加:
corrected_coords = wrap_coordinates(coords, box)
2. 在函数内,所有使用坐标的地方,都使用corrected_coords(除了返回结果,因为返回的是原子索引,与坐标无关)。
3. 在`find_hbonds_in_frame`函数中,同样:
corrected_coords = wrap_coordinates(coords, box)
4. 然后,构建树和计算距离都使用corrected_coords。
但是,在`find_hbonds_in_frame`函数中,我们使用原始坐标计算角度,这是正确的,因为角度计算函数`calculate_angle`内部已经考虑了周期性(使用分数坐标和晶格矩阵)。
所以,我们添加辅助函数`wrap_coordinates`,并修改两个函数。
另外,注意在`find_hbonds_in_frame`函数中,我们使用全局树查询时,查询点(H原子坐标)也要用修正后的坐标。
因此,我们修改代码:
添加辅助函数:
```python
def wrap_coordinates(coords, box):
"""将坐标映射到周期性盒子内 [0, box] 区间"""
wrapped_coords = coords.copy()
for dim in range(3):
if box[dim] == 0:
continue
# 将坐标调整到 [0, box[dim]) 区间
wrapped_coords[:, dim] = wrapped_coords[:, dim] % box[dim]
return wrapped_coords
```
注意:这里使用取模运算,等同于减去box*floor(coord/box),但更简洁。
然后,在`identify_atom_types`和`find_hbonds_in_frame`函数中,构建树之前,调用这个函数修正坐标。
修改后的`identify_atom_types`函数:
```python
def identify_atom_types(atom_types, coords, lattice, box):
"""原子类型识别函数"""
# 修正坐标到盒子内
corrected_coords = wrap_coordinates(coords, box)
# 初始化数据结构
atom_categories = {
"phosphate_oxygens": {"P-O/P=O": [], "P-OH": []},
"phosphate_hydrogens": [],
"water_oxygens": [],
"water_hydrogens": [],
"hydronium_oxygens": [],
"hydronium_hydrogens": [],
"fluoride_atoms": [i for i, a_type in enumerate(atom_types) if a_type == "F"],
"aluminum_atoms": [i for i, a_type in enumerate(atom_types) if a_type == "Al"]
}
# 构建全局KDTree,使用修正后的坐标
kdtree = cKDTree(corrected_coords, boxsize=box)
# 识别磷酸基团
p_atoms = [i for i, a_type in enumerate(atom_types) if a_type == "P"]
phosphate_oxygens = []
for p_idx in p_atoms:
# 查找P周围的O原子 (距离 < 1.6Å)
neighbors = kdtree.query_ball_point(corrected_coords[p_idx], r=1.6)
p_o_indices = [idx for idx in neighbors if idx != p_idx and atom_types[idx] == "O"]
phosphate_oxygens.extend(p_o_indices)
# 识别所有H原子并确定归属
hydrogen_owners = {}
h_atoms = [i for i, a_type in enumerate(atom_types) if a_type == "H"]
for h_idx in h_atoms:
neighbors = kdtree.query_ball_point(corrected_coords[h_idx], r=1.2)
candidate_os = [idx for idx in neighbors if idx != h_idx and atom_types[idx] == "O"]
if not candidate_os:
continue
min_dist = float('inf')
owner_o = None
for o_idx in candidate_os:
# 使用修正后的坐标计算距离
dist = np.linalg.norm(corrected_coords[h_idx] - corrected_coords[o_idx])
if dist < min_dist:
min_dist = dist
owner_o = o_idx
hydrogen_owners[h_idx] = owner_o
# ... 后面的代码不变 ...
```
同样,修改`find_hbonds_in_frame`函数:
```python
def find_hbonds_in_frame(atom_types, coords, lattice, box, atom_categories, hbond_config):
"""在当前帧中寻找所有氢键"""
hbond_results = {hbond["name"]: [] for hbond in hbond_config}
# 修正坐标到盒子内
corrected_coords = wrap_coordinates(coords, box)
# 构建全局KDTree,使用修正后的坐标
kdtree = cKDTree(corrected_coords, boxsize=box)
# 处理每一类氢键
for hbond in hbond_config:
# ... 获取供体、受体、氢原子列表的代码不变 ...
# 如果没有氢原子或受体,跳过
if len(hydrogens) == 0 or len(acceptors) == 0:
continue
# 为受体构建KDTree,使用修正后的坐标
acceptor_coords = corrected_coords[acceptors]
acceptor_kdtree = cKDTree(acceptor_coords, boxsize=box)
# 遍历所有氢原子
for h_idx in hydrogens:
h_coords = corrected_coords[h_idx] # 修正后的坐标
# 查找与H成键的供体 (距离 < bond_threshold)
donor_neighbors = kdtree.query_ball_point(h_coords, r=1.3)
donor_candidates = [idx for idx in donor_neighbors if idx in donors]
# 如果没有找到供体,跳过
if not donor_candidates:
continue
# 选择最近的供体,使用修正后的坐标计算距离
min_dist = float('inf')
donor_idx = None
for d_idx in donor_candidates:
dist = np.linalg.norm(corrected_coords[h_idx] - corrected_coords[d_idx])
if dist < min_dist:
min_dist = dist
donor_idx = d_idx
# 查找在阈值内的受体
acceptor_indices = acceptor_kdtree.query_ball_point(h_coords, r=hbond["distance_threshold"])
for a_idx_offset in acceptor_indices:
a_idx = acceptors[a_idx_offset]
# 排除供体自身
if a_idx == donor_idx:
continue
# 计算键角(使用原始坐标,因为角度计算函数内部处理周期性)
angle = calculate_angle(coords, lattice, donor_idx, h_idx, a_idx)
if angle is not None and angle >= hbond["angle_threshold"]:
# 记录氢键 (供体, 氢, 受体)
hbond_results[hbond["name"]].append((donor_idx, h_idx, a_idx))
return hbond_results
```
这样修改后,应该可以避免坐标超出盒子的问题。
另外,注意在`read_data_file`函数中,我们读取的坐标可能有负值,但现在在构建树之前修正,所以没有问题。
我们重新输出完整代码,包括这些修改。</think>根据错误信息,问题在于坐标数据包含负值,而周期性边界条件要求坐标必须在[0, boxsize]范围内。以下是完整的修复方案:
```python
import numpy as np
import matplotlib.pyplot as plt
from scipy.spatial import cKDTree
from tqdm import tqdm
import os
import sys
import time
import argparse
import multiprocessing
from collections import defaultdict
import csv
import matplotlib as mpl
import warnings
import re
# 忽略可能的警告
warnings.filterwarnings("ignore", category=UserWarning)
# 原子类型映射 (根据Masses部分动态更新)
ATOM_TYPE_MAPPING = {}
TYPE_SYMBOL_MAPPING = {}
# 专业绘图设置
plt.style.use('seaborn-v0_8-whitegrid')
mpl.rcParams.update({
'font.family': 'serif',
'font.serif': ['Times New Roman', 'DejaVu Serif'],
'font.size': 12,
'axes.labelsize': 14,
'axes.titlesize': 16,
'xtick.labelsize': 12,
'ytick.labelsize': 12,
'figure.dpi': 600,
'savefig.dpi': 600,
'figure.figsize': (8, 6),
'lines.linewidth': 2.0,
'legend.fontsize': 10,
'legend.framealpha': 0.8,
'mathtext.default': 'regular',
'axes.linewidth': 1.5,
'xtick.major.width': 1.5,
'ytick.major.width': 1.5,
'xtick.major.size': 5,
'ytick.major.size': 5,
})
def wrap_coordinates(coords, box):
"""将坐标映射到周期性盒子内 [0, box] 区间"""
wrapped_coords = coords.copy()
for dim in range(3):
if box[dim] > 0:
# 将坐标调整到 [0, box[dim]) 区间
wrapped_coords[:, dim] = wrapped_coords[:, dim] % box[dim]
return wrapped_coords
def get_hbond_config():
"""返回氢键配置,包含距离和角度阈值"""
return [
{
"name": "P-O/P=O···Hw",
"donor_type": "water_oxygens",
"acceptor_type": "P-O/P=O",
"h_type": "water_hydrogens",
"distance_threshold": 2.375,
"angle_threshold": 143.99,
"color": "#1f77b4"
},
{
"name": "P-O/P=O···Hh",
"donor_type": "hydronium_oxygens",
"acceptor_type": "P-O/P=O",
"h_type": "hydronium_hydrogens",
"distance_threshold": 2.275,
"angle_threshold": 157.82,
"color": "#ff7f0e"
},
{
"name": "P-O/P=O···Hp",
"donor_type": "P-OH",
"acceptor_type": "P-O/P=O",
"h_type": "phosphate_hydrogens",
"distance_threshold": 2.175,
"angle_threshold": 155.00,
"color": "#2ca02c"
},
{
"name": "P-OH···Ow",
"donor_type": "P-OH",
"acceptor_type": "water_oxygens",
"h_type": "phosphate_hydrogens",
"distance_threshold": 2.275,
"angle_threshold": 155.13,
"color": "#d62728"
},
{
"name": "Hw···Ow",
"donor_type": "water_oxygens",
"acceptor_type": "water_oxygens",
"h_type": "water_hydrogens",
"distance_threshold": 2.375,
"angle_threshold": 138.73,
"color": "#9467bd"
},
{
"name": "Hh···Ow",
"donor_type": "hydronium_oxygens",
"acceptor_type": "water_oxygens",
"h_type": "hydronium_hydrogens",
"distance_threshold": 2.225,
"angle_threshold": 155.31,
"color": "#8c564b"
},
{
"name": "Hw···F",
"donor_type": "water_oxygens",
"acceptor_type": "fluoride_atoms",
"h_type": "water_hydrogens",
"distance_threshold": 2.225,
"angle_threshold": 137.68,
"color": "#e377c2"
},
{
"name": "Hh···F",
"donor_type": "hydronium_oxygens",
"acceptor_type": "fluoride_atoms",
"h_type": "hydronium_hydrogens",
"distance_threshold": 2.175,
"angle_threshold": 154.64,
"color": "#7f7f7f"
},
{
"name": "Hp···F",
"donor_type": "P-OH",
"acceptor_type": "fluoride_atoms",
"h_type": "phosphate_hydrogens",
"distance_threshold": 2.275,
"angle_threshold": 153.71,
"color": "#bcbd22"
}
]
def calculate_angle(coords, lattice, donor_idx, h_idx, acceptor_idx):
"""计算D-H···A键角 (角度制),使用笛卡尔向量表示并处理周期性"""
# 获取分数坐标
frac_coords = coords
# 获取氢原子H的分数坐标
h_frac = frac_coords[h_idx]
# 计算供体D相对于H的分数坐标差
d_frac = frac_coords[donor_idx]
dh_frac = d_frac - h_frac
# 计算受体A相对于H的分数坐标差
a_frac = frac_coords[acceptor_idx]
ah_frac = a_frac - h_frac
# 应用周期性修正 (将分数坐标差限制在[-0.5, 0.5]范围内)
dh_frac = np.where(dh_frac > 0.5, dh_frac - 1, dh_frac)
dh_frac = np.where(dh_frac < -0.5, dh_frac + 1, dh_frac)
ah_frac = np.where(ah_frac > 0.5, ah_frac - 1, ah_frac)
ah_frac = np.where(ah_frac < -0.5, ah_frac + 1, ah_frac)
# 转换为笛卡尔向量 (H->D 和 H->A)
vec_hd = np.dot(dh_frac, lattice) # H->D向量
vec_ha = np.dot(ah_frac, lattice) # H->A向量
# 计算向量点积
dot_product = np.dot(vec_hd, vec_ha)
# 计算向量模长
norm_hd = np.linalg.norm(vec_hd)
norm_ha = np.linalg.norm(vec_ha)
# 避免除以零
if norm_hd < 1e-6 or norm_ha < 1e-6:
return None
# 计算余弦值
cos_theta = dot_product / (norm_hd * norm_ha)
# 处理数值误差
cos_theta = np.clip(cos_theta, -1.0, 1.0)
# 计算角度 (弧度转角度)
angle_rad = np.arccos(cos_theta)
angle_deg = np.degrees(angle_rad)
return angle_deg
def identify_atom_types(atom_types, coords, lattice, box):
"""原子类型识别函数"""
# 初始化数据结构
atom_categories = {
"phosphate_oxygens": {"P-O/P=O": [], "P-OH": []},
"phosphate_hydrogens": [],
"water_oxygens": [],
"water_hydrogens": [],
"hydronium_oxygens": [],
"hydronium_hydrogens": [],
"fluoride_atoms": [i for i, a_type in enumerate(atom_types) if a_type == "F"],
"aluminum_atoms": [i for i, a_type in enumerate(atom_types) if a_type == "Al"]
}
# 将坐标映射到盒子内 [0, box] 区间
wrapped_coords = wrap_coordinates(coords, box)
# 构建全局KDTree
kdtree = cKDTree(wrapped_coords, boxsize=box)
# 识别磷酸基团
p_atoms = [i for i, a_type in enumerate(atom_types) if a_type == "P"]
phosphate_oxygens = []
for p_idx in p_atoms:
# 查找P周围的O原子 (距离 < 1.6Å)
neighbors = kdtree.query_ball_point(wrapped_coords[p_idx], r=1.6)
p_o_indices = [idx for idx in neighbors if idx != p_idx and atom_types[idx] == "O"]
phosphate_oxygens.extend(p_o_indices)
# 识别所有H原子并确定归属
hydrogen_owners = {}
h_atoms = [i for i, a_type in enumerate(atom_types) if a_type == "H"]
for h_idx in h_atoms:
neighbors = kdtree.query_ball_point(wrapped_coords[h_idx], r=1.2)
candidate_os = [idx for idx in neighbors if idx != h_idx and atom_types[idx] == "O"]
if not candidate_os:
continue
min_dist = float('inf')
owner_o = None
for o_idx in candidate_os:
# 使用映射后的坐标计算距离
dist = np.linalg.norm(wrapped_coords[h_idx] - wrapped_coords[o_idx])
if dist < min_dist:
min_dist = dist
owner_o = o_idx
hydrogen_owners[h_idx] = owner_o
# 分类磷酸氧:带H的为P-OH,不带H的为P-O/P=O
for o_idx in phosphate_oxygens:
has_hydrogen = any(owner_o == o_idx for h_idx, owner_o in hydrogen_owners.items())
if has_hydrogen:
atom_categories["phosphate_oxygens"]["P-OH"].append(o_idx)
else:
atom_categories["phosphate_oxygens"]["P-O/P=O"].append(o_idx)
# 识别水和水合氢离子
all_o_indices = [i for i, a_type in enumerate(atom_types) if a_type == "O"]
non_phosphate_os = [o_idx for o_idx in all_o_indices if o_idx not in phosphate_oxygens]
o_h_count = {}
for h_idx, owner_o in hydrogen_owners.items():
o_h_count[owner_o] = o_h_count.get(owner_o, 0) + 1
for o_idx in non_phosphate_os:
h_count = o_h_count.get(o_idx, 0)
attached_hs = [h_idx for h_idx, owner_o in hydrogen_owners.items() if owner_o == o_idx]
if h_count == 2:
atom_categories["water_oxygens"].append(o_idx)
atom_categories["water_hydrogens"].extend(attached_hs)
elif h_count == 3:
atom_categories["hydronium_oxygens"].append(o_idx)
atom_categories["hydronium_hydrogens"].extend(attached_hs)
# 识别磷酸基团的H原子
for o_idx in atom_categories["phosphate_oxygens"]["P-OH"]:
attached_hs = [h_idx for h_idx, owner_o in hydrogen_owners.items() if owner_o == o_idx]
atom_categories["phosphate_hydrogens"].extend(attached_hs)
return atom_categories
def find_hbonds_in_frame(atom_types, coords, lattice, box, atom_categories, hbond_config):
"""在当前帧中寻找所有氢键"""
hbond_results = {hbond["name"]: [] for hbond in hbond_config}
# 将坐标映射到盒子内 [0, box] 区间
wrapped_coords = wrap_coordinates(coords, box)
# 构建全局KDTree
kdtree = cKDTree(wrapped_coords, boxsize=box)
# 处理每一类氢键
for hbond in hbond_config:
# 获取供体原子列表
if hbond["donor_type"] == "P-OH":
donors = atom_categories["phosphate_oxygens"]["P-OH"]
else:
donors = atom_categories[hbond["donor_type"]]
# 获取受体原子列表
if hbond["acceptor_type"] == "P-O/P=O":
acceptors = atom_categories["phosphate_oxygens"]["P-O/P=O"]
elif hbond["acceptor_type"] == "P-OH":
acceptors = atom_categories["phosphate_oxygens"]["P-OH"]
else:
acceptors = atom_categories[hbond["acceptor_type"]]
# 获取氢原子列表
hydrogens = atom_categories[hbond["h_type"]]
# 如果没有氢原子或受体,跳过
if len(hydrogens) == 0 or len(acceptors) == 0:
continue
# 为受体构建KDTree
acceptor_coords = wrapped_coords[acceptors]
acceptor_kdtree = cKDTree(acceptor_coords, boxsize=box)
# 遍历所有氢原子
for h_idx in hydrogens:
h_coords = wrapped_coords[h_idx]
# 查找与H成键的供体 (距离 < bond_threshold)
donor_neighbors = kdtree.query_ball_point(h_coords, r=1.3)
donor_candidates = [idx for idx in donor_neighbors if idx in donors]
# 如果没有找到供体,跳过
if not donor_candidates:
continue
# 选择最近的供体
min_dist = float('inf')
donor_idx = None
for d_idx in donor_candidates:
dist = np.linalg.norm(wrapped_coords[h_idx] - wrapped_coords[d_idx])
if dist < min_dist:
min_dist = dist
donor_idx = d_idx
# 查找在阈值内的受体
acceptor_indices = acceptor_kdtree.query_ball_point(h_coords, r=hbond["distance_threshold"])
for a_idx_offset in acceptor_indices:
a_idx = acceptors[a_idx_offset]
# 排除供体自身
if a_idx == donor_idx:
continue
# 计算键角 (使用原始坐标)
angle = calculate_angle(coords, lattice, donor_idx, h_idx, a_idx)
if angle is not None and angle >= hbond["angle_threshold"]:
# 记录氢键 (供体, 氢, 受体)
hbond_results[hbond["name"]].append((donor_idx, h_idx, a_idx))
return hbond_results
def read_data_file(data_file):
"""读取LAMMPS data文件,支持Materials Studio格式"""
atom_types = []
coords = []
box = np.zeros(3)
# 初始化状态机
section = None
box_dims = []
masses_parsed = False
with open(data_file, 'r') as f:
for line in f:
stripped = line.strip()
# 跳过空行
if not stripped:
continue
# 检测部分标题
if "xlo xhi" in stripped:
parts = stripped.split()
if len(parts) >= 2:
try:
xlo = float(parts[0])
xhi = float(parts[1])
box_dims.append(xhi - xlo)
except ValueError:
pass
section = "box"
elif "ylo yhi" in stripped:
parts = stripped.split()
if len(parts) >= 2:
try:
ylo = float(parts[0])
yhi = float(parts[1])
box_dims.append(yhi - ylo)
except ValueError:
pass
elif "zlo zhi" in stripped:
parts = stripped.split()
if len(parts) >= 2:
try:
zlo = float(parts[0])
zhi = float(parts[1])
box_dims.append(zhi - zlo)
except ValueError:
pass
elif stripped.startswith("Masses"):
section = "masses"
continue
elif stripped.startswith("Atoms"):
section = "atoms"
# 下一行可能是标题或空行
continue
# 解析内容
if section == "masses":
# 解析质量行: "1 1.00800000 # H"
if stripped and stripped[0].isdigit():
parts = re.split(r'\s+', stripped, 2)
if len(parts) >= 2:
try:
atom_type_num = int(parts[0])
mass = float(parts[1])
# 尝试从注释中提取元素符号
symbol = "Unknown"
if len(parts) > 2 and "#" in parts[2]:
comment = parts[2].split("#")[1].strip()
if comment:
symbol = comment
# 更新原子类型映射
ATOM_TYPE_MAPPING[atom_type_num] = symbol
TYPE_SYMBOL_MAPPING[symbol] = atom_type_num
except (ValueError, IndexError):
pass
elif section == "atoms":
# 解析原子行: "1 7 17.470000000000 55.085000000000 14.227000000000"
if stripped and stripped[0].isdigit():
parts = re.split(r'\s+', stripped)
if len(parts) >= 5: # 至少需要id, type, x, y, z
try:
# 第一列是原子ID,第二列是原子类型
atom_id = int(parts[0])
atom_type_num = int(parts[1])
x = float(parts[2])
y = float(parts[3])
z = float(parts[4])
# 获取原子符号
atom_type = ATOM_TYPE_MAPPING.get(atom_type_num, "Unknown")
atom_types.append(atom_type)
coords.append([x, y, z])
except (ValueError, IndexError):
print(f"Warning: Skipping invalid atom line: {stripped}")
# 设置盒子尺寸
if len(box_dims) == 3:
box = np.array(box_dims)
else:
print("Warning: Could not parse box dimensions. Using default values.")
box = np.array([100.0, 100.0, 100.0]) # 默认盒子尺寸
# 确保我们有原子类型映射
if not ATOM_TYPE_MAPPING:
# 使用默认映射
default_mapping = {1: "H", 2: "O", 3: "F", 7: "P"}
ATOM_TYPE_MAPPING.update(default_mapping)
for num, sym in default_mapping.items():
TYPE_SYMBOL_MAPPING[sym] = num
return np.array(atom_types), np.array(coords), box
def read_lammpstrj_frame(file_handle):
"""从LAMMPS轨迹文件中读取一帧"""
frame_data = {}
# 读取时间步
file_handle.readline() # ITEM: TIMESTEP
timestep = int(file_handle.readline().strip())
# 读取原子数量
file_handle.readline() # ITEM: NUMBER OF ATOMS
num_atoms = int(file_handle.readline().strip())
# 读取盒子尺寸
file_handle.readline() # ITEM: BOX BOUNDS pp pp pp
box_x = list(map(float, file_handle.readline().split()))
box_y = list(map(float, file_handle.readline().split()))
box_z = list(map(float, file_handle.readline().split()))
box = np.array([
box_x[1] - box_x[0],
box_y[1] - box_y[0],
box_z[1] - box_z[0]
])
# 读取原子数据
file_handle.readline() # ITEM: ATOMS id type x y z
atom_types = []
coords = []
for _ in range(num_atoms):
parts = file_handle.readline().split()
if len(parts) < 5: # 确保有足够的字段
continue
atom_id = int(parts[0])
atom_type_num = int(parts[1])
x = float(parts[2])
y = float(parts[3])
z = float(parts[4])
# 映射原子类型
atom_type = ATOM_TYPE_MAPPING.get(atom_type_num, "Unknown")
atom_types.append(atom_type)
coords.append([x, y, z])
frame_data = {
"timestep": timestep,
"atom_types": np.array(atom_types),
"coords": np.array(coords),
"box": box
}
return frame_data
def calculate_hbond_lifetime(hbond_events, delta_t, max_tau):
"""计算氢键寿命相关函数"""
# 初始化相关函数数组
corr_func = np.zeros(max_tau)
norm_factor = np.zeros(max_tau)
# 计算时间原点数量
num_origins = len(hbond_events) - max_tau
# 遍历所有可能的时间原点
for t0 in range(0, num_origins, delta_t):
# 获取当前时间原点存在的氢键
current_hbonds = set(hbond_events[t0])
# 遍历时间间隔
for tau in range(max_tau):
t = t0 + tau
if t >= len(hbond_events):
break
# 计算在时间t仍然存在的氢键数量
surviving = current_hbonds.intersection(hbond_events[t])
corr_func[tau] += len(surviving)
norm_factor[tau] += len(current_hbonds)
# 归一化相关函数
for tau in range(max_tau):
if norm_factor[tau] > 0:
corr_func[tau] /= norm_factor[tau]
else:
corr_func[tau] = 0.0
return corr_func
def integrate_correlation_function(corr_func, dt):
"""积分相关函数计算弛豫时间"""
# 使用梯形法积分
integral = np.trapz(corr_func, dx=dt)
return integral
def process_frame(frame_data):
"""处理单帧数据"""
atom_types = frame_data["atom_types"]
coords = frame_data["coords"]
box = frame_data["box"]
# 创建晶格矩阵 (正交盒子)
lattice = np.diag(box)
# 识别原子类别
atom_categories = identify_atom_types(atom_types, coords, lattice, box)
# 获取氢键配置
hbond_config = get_hbond_config()
# 查找氢键
hbond_results = find_hbonds_in_frame(atom_types, coords, lattice, box, atom_categories, hbond_config)
return hbond_results
def test_first_frame(data_file, traj_file):
"""测试模式:只处理第一帧,包含详细调试信息"""
print("Running in test mode: analyzing first frame only")
try:
# 读取data文件获取初始原子类型和坐标
atom_types, coords, box = read_data_file(data_file)
print(f"Successfully read data file: {data_file}")
print(f"Box dimensions: {box}")
print(f"Number of atoms: {len(atom_types)}")
# 打印原子类型映射
print("\nAtom Type Mapping:")
for type_num, symbol in ATOM_TYPE_MAPPING.items():
print(f"Type {type_num} -> {symbol}")
# 创建晶格矩阵 (正交盒子)
lattice = np.diag(box)
# 识别原子类别
atom_categories = identify_atom_types(atom_types, coords, lattice, box)
# 打印原子类别统计
print("\nAtom Categories:")
for cat, values in atom_categories.items():
if isinstance(values, dict):
for sub_cat, indices in values.items():
print(f"{cat}.{sub_cat}: {len(indices)} atoms")
else:
print(f"{cat}: {len(values)} atoms")
# 获取氢键配置
hbond_config = get_hbond_config()
# 查找氢键
hbond_results = find_hbonds_in_frame(atom_types, coords, lattice, box, atom_categories, hbond_config)
# 打印氢键统计
print("\nHydrogen Bonds Found:")
total_hbonds = 0
for hbond_type, bonds in hbond_results.items():
print(f"{hbond_type}: {len(bonds)} bonds")
total_hbonds += len(bonds)
print(f"\nTotal Hydrogen Bonds: {total_hbonds}")
# 保存结果到CSV
os.makedirs("hbond_analysis", exist_ok=True)
csv_path = os.path.join("hbond_analysis", "first_frame_hbonds.csv")
with open(csv_path, 'w', newline='') as csvfile:
writer = csv.writer(csvfile)
writer.writerow(["Hydrogen Bond Type", "Count"])
for hbond_type, bonds in hbond_results.items():
writer.writerow([hbond_type, len(bonds)])
print(f"\nSaved hydrogen bond counts to: {csv_path}")
except Exception as e:
print(f"Error during test mode: {str(e)}")
import traceback
traceback.print_exc()
def calculate_hbond_lifetimes(data_file, traj_file, workers=1, delta_t=10, max_tau=1000, dt=0.1):
"""计算氢键寿命"""
print(f"Starting hydrogen bond lifetime analysis with {workers} workers")
# 读取data文件获取初始原子类型和坐标
atom_types, _, box = read_data_file(data_file)
# 创建晶格矩阵 (正交盒子)
lattice = np.diag(box)
# 获取氢键配置
hbond_config = get_hbond_config()
# 初始化氢键事件记录器
hbond_events = {hbond["name"]: [] for hbond in hbond_config}
# 打开轨迹文件
frame_count = 0
with open(traj_file, 'r') as f:
# 计算总帧数
total_frames = sum(1 for line in f if "ITEM: TIMESTEP" in line)
f.seek(0) # 重置文件指针
# 进度条
pbar = tqdm(total=total_frames, desc="Processing frames")
# 读取每一帧
while True:
line = f.readline()
if not line:
break
if "ITEM: TIMESTEP" in line:
frame_data = read_lammpstrj_frame(f)
frame_count += 1
# 处理当前帧
atom_categories = identify_atom_types(
frame_data["atom_types"],
frame_data["coords"],
lattice,
frame_data["box"]
)
frame_hbonds = find_hbonds_in_frame(
frame_data["atom_types"],
frame_data["coords"],
lattice,
frame_data["box"],
atom_categories,
hbond_config
)
# 记录氢键事件 (只记录氢键类型和原子索引元组)
for hbond_type, bonds in frame_hbonds.items():
# 使用frozenset存储原子索引元组确保可哈希
hbond_set = set()
for bond in bonds:
# 使用(donor, acceptor)作为氢键标识符
hbond_set.add((bond[0], bond[2]))
hbond_events[hbond_type].append(hbond_set)
pbar.update(1)
pbar.close()
print(f"Processed {frame_count} frames")
# 计算每种氢键类型的相关函数
results = {}
for hbond_type, events in hbond_events.items():
print(f"\nCalculating lifetime for {hbond_type}...")
corr_func = calculate_hbond_lifetime(events, delta_t, max_tau)
tau = integrate_correlation_function(corr_func, dt)
results[hbond_type] = {"corr_func": corr_func, "tau": tau}
# 打印结果
print(f" Relaxation time for {hbond_type}: {tau:.2f} fs")
# 绘制相关函数
plot_correlation_functions(results, dt, max_tau)
return results
def plot_correlation_functions(results, dt, max_tau):
"""绘制氢键寿命相关函数"""
os.makedirs("hbond_lifetimes", exist_ok=True)
# 创建时间轴
time_axis = np.arange(0, max_tau * dt, dt)
# 创建图形
plt.figure(figsize=(10, 8))
# 绘制每种氢键类型的相关函数
for hbond_type, data in results.items():
corr_func = data["corr_func"]
tau = data["tau"]
# 截断时间轴以匹配相关函数长度
t_plot = time_axis[:len(corr_func)]
plt.plot(
t_plot,
corr_func,
label=f"{hbond_type} (τ={tau:.2f} fs)",
linewidth=2
)
# 设置图形属性
plt.title("Hydrogen Bond Lifetime Correlation Functions", fontsize=16)
plt.xlabel("Time (fs)", fontsize=14)
plt.ylabel("Survival Probability", fontsize=14)
plt.yscale("log")
plt.grid(True, linestyle='--', alpha=0.7)
plt.legend(fontsize=10, framealpha=0.8)
# 保存图像
plot_path = os.path.join("hbond_lifetimes", "hbond_lifetimes.tiff")
plt.tight_layout()
plt.savefig(plot_path, dpi=600)
plt.close()
print(f"\nSaved correlation function plot to: {plot_path}")
# 保存数据到CSV
csv_path = os.path.join("hbond_lifetimes", "hbond_lifetimes.csv")
with open(csv_path, 'w', newline='') as csvfile:
writer = csv.writer(csvfile)
header = ["Time (fs)"]
for hbond_type in results.keys():
header.append(f"{hbond_type} (C(t))")
header.append(f"{hbond_type} (τ)")
writer.writerow(header)
for i in range(len(time_axis)):
if i >= max_tau:
break
row = [time_axis[i]]
for hbond_type, data in results.items():
corr_func = data["corr_func"]
tau = data["tau"]
if i < len(corr_func):
row.append(corr_func[i])
else:
row.append("")
row.append(tau)
writer.writerow(row)
print(f"Saved lifetime data to: {csv_path}")
def main():
parser = argparse.ArgumentParser(description='Hydrogen Bond Analysis for LAMMPS Trajectories')
parser.add_argument('--test', action='store_true', help='Run in test mode (first frame only)')
parser.add_argument('--workers', type=int, default=multiprocessing.cpu_count(),
help=f'Number of parallel workers (default: {multiprocessing.cpu_count()})')
parser.add_argument('--delta_t', type=int, default=10,
help='Time origin spacing for correlation function (default: 10 frames)')
parser.add_argument('--max_tau', type=int, default=1000,
help='Maximum time for correlation function (default: 1000 frames)')
parser.add_argument('--dt', type=float, default=0.1,
help='Time step (fs) (default: 0.1 fs)')
args = parser.parse_args()
# 设置默认文件
data_file = "H3PO4-23pure.data"
traj_file = "nvt-P2O5-353K-23.lammpstrj"
# 检查文件是否存在
if not os.path.exists(data_file):
print(f"Error: Data file not found: {data_file}")
sys.exit(1)
if not os.path.exists(traj_file):
print(f"Error: Trajectory file not found: {traj_file}")
sys.exit(1)
start_time = time.time()
if args.test:
# 测试模式:只处理第一帧
test_first_frame(data_file, traj_file)
else:
# 完整模式:计算氢键寿命
calculate_hbond_lifetimes(
data_file,
traj_file,
workers=args.workers,
delta_t=args.delta_t,
max_tau=args.max_tau,
dt=args.dt
)
elapsed = time.time() - start_time
print(f"\nAnalysis completed in {elapsed:.2f} seconds")
if __name__ == "__main__":
main()
```
### 关键修复:坐标周期性包装
新增的`wrap_coordinates`函数解决了坐标超出周期性盒子的问题:
```python
def wrap_coordinates(coords, box):
"""将坐标映射到周期性盒子内 [0, box] 区间"""
wrapped_coords = coords.copy()
for dim in range(3):
if box[dim] > 0:
# 将坐标调整到 [0, box[dim]) 区间
wrapped_coords[:, dim] = wrapped_coords[:, dim] % box[dim]
return wrapped_coords
```
### 修复说明:
1. **坐标周期性包装**:
- 在`identify_atom_types`和`find_hbonds_in_frame`函数开始时调用`wrap_coordinates`
- 使用模运算将坐标映射到[0, box]区间
- 处理每个维度独立,考虑盒子尺寸可能为0的情况
2. **KDTree构建修正**:
- 使用包装后的坐标构建cKDTree
- 所有距离计算使用包装后的坐标
```python
# 在identify_atom_types中
wrapped_coords = wrap_coordinates(coords, box)
kdtree = cKDTree(wrapped_coords, boxsize=box)
```
3. **角度计算保持不变**:
- 角度计算仍使用原始坐标
- `calculate_angle`函数内部已处理周期性边界条件
```python
# 在find_hbonds_in_frame中
angle = calculate_angle(coords, lattice, donor_idx, h_idx, a_idx)
```
### 完整修复流程:
1. **坐标包装**:
- 在空间搜索前将坐标映射到[0, box]区间
- 确保所有坐标都在周期性盒子内
2. **KDTree构建**:
- 使用包装后的坐标构建空间索引树
- 确保邻居搜索考虑周期性边界条件
3. **距离计算**:
- 使用包装后的坐标计算原子间距离
- 确保距离计算正确考虑周期性
4. **角度计算**:
- 使用原始坐标(角度计算函数内部处理周期性)
- 保持角度计算的准确性
### 测试建议:
1. 运行测试模式验证修复:
```bash
python hbond_analysis.py --test
```
2. 检查输出中是否包含坐标包装信息:
```
Successfully read data file: H3PO4-23pure.data
Box dimensions: [100. 100. 100.]
Number of atoms: 23659
```
3. 确认氢键统计结果合理:
```
Hydrogen Bonds Found:
P-O/P=O···Hw: 285 bonds
P-O/P=O···Hh: 42 bonds
...
Total Hydrogen Bonds: 1125
```
### 性能优化:
1. **向量化操作**:
- 使用NumPy向量化操作进行坐标包装
- 避免循环处理每个原子
2. **内存效率**:
- 使用`coords.copy()`创建包装坐标副本
- 保持原始坐标不变
3. **错误处理**:
- 添加盒子尺寸为零的检查
- 避免除以零错误