匹配特定正则模式统计原子行数并添加至列表

需求说明 #

读取LAMMPS输出文件，在每个ITEM: ATOMS id mol element type x y z头部到下一个ITEM: TIMESTEP头部的区间内，统计符合以下正则表达式的行的数量：

\d+\s\d+\s[N]\s[1]\s\d+\.\d+\s\d+\.\d+\s\d+\.\d+

将每个区间的统计结果依次存入列表。

示例文件 #

ITEM: TIMESTEP
0
ITEM: NUMBER OF ATOMS
38873
ITEM: BOX BOUNDS pp pp pp
5.7248792664789718e+00 8.8775120733336294e+01
5.7248792664789718e+00 8.8775120733336294e+01
5.7248792664789718e+00 8.8775120733336294e+01
ITEM: ATOMS id mol element type x y z
1 1 N 1 36.8189 69.7138 60.6542
2 1 H 2 37.1158 70.7384 60.4896
3 1 H 2 36.2448 69.5936 61.4754
4 1 C 3 37.692 68.7914 60.2942
5 1 C 4 37.5524 67.3932 60.6119
ITEM: TIMESTEP
0
ITEM: NUMBER OF ATOMS
38873
ITEM: BOX BOUNDS pp pp pp
5.7248792664789718e+00 8.8775120733336294e+01
5.7248792664789718e+00 8.8775120733336294e+01
5.7248792664789718e+00 8.8775120733336294e+01
ITEM: ATOMS id mol element type x y z
1 1 N 1 36.8189 69.7138 60.6542
2 1 H 2 37.1158 70.7384 60.4896
3 1 N 1 36.2448 69.5936 61.4754
4 1 N 1 37.692 68.7914 60.2942
5 1 C 4 37.5524 67.3932 60.6119

预期结果 #

N_list = [1, 3]

Python实现代码 #

import re

def count_target_atoms(file_path):
    N_list = []
    in_atoms_section = False
    current_count = 0
    # 编译正则表达式，匹配指定格式的原子行
    pattern = re.compile(r"\d+\s\d+\s[N]\s[1]\s\d+\.\d+\s\d+\.\d+\s\d+\.\d+")
    
    with open(file_path, 'r') as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            # 进入ATOMS数据区间，初始化计数器
            if line == "ITEM: ATOMS id mol element type x y z":
                in_atoms_section = True
                current_count = 0
                continue
            # 离开ATOMS区间，保存当前计数
            if line == "ITEM: TIMESTEP" and in_atoms_section:
                in_atoms_section = False
                N_list.append(current_count)
                continue
            # 在ATOMS区间内，匹配目标行并计数
            if in_atoms_section:
                if pattern.match(line):
                    current_count += 1
        # 处理文件末尾无ITEM: TIMESTEP的情况，保存最后一个区间的计数
        if in_atoms_section:
            N_list.append(current_count)
    return N_list

# 使用示例，替换为你的文件路径
file_path = "lammps_output.txt"
N_list = count_target_atoms(file_path)
print(f"N_list = {N_list}")

代码说明 #

• 用in_atoms_section标记是否处于ATOMS数据区间
• 遇到ATOMS头部时重置计数器，进入统计状态
• 遇到TIMESTEP头部且处于统计状态时，将当前计数存入列表并退出统计
• 区间内每匹配一行目标正则，计数器加1
• 兼容文件结尾没有TIMESTEP头部的情况，确保最后一个区间的计数被保存

内容的提问来源于stack exchange，提问作者kata248