split_set 模块说明
概述
split_set.py 模块是 PhyMLP Kit 工具包中的关键数据处理工具，专门用于将 EXTXYZ 格式的原子结构数据集分割为训练集和验证集。该模块是机器学习势训练工作流中数据准备阶段的核心组件。
功能特性
特性 |
描述 |
|---|---|
随机分割 |
完全随机打乱数据集，确保统计无偏性 |
比例可控 |
可精确控制验证集比例（0.0-1.0） |
格式保持 |
保持 EXTXYZ 格式完整，不丢失任何元数据 |
高效处理 |
支持处理包含数万结构的大型数据集 |
简单易用 |
提供命令行和 Python API 两种使用方式 |
快速开始
基本使用方式
# 方式1：通过 phymlp-kit 命令行
phymlp-kit split_set --input_file full.extxyz --val_ratio 0.2
# 方式2：直接运行 Python 脚本
python split_set.py --input full.extxyz --val_ratio 0.15
# 方式3：在 Python 代码中导入
from phymlp.tools.split_set import split_extxyz_structures
split_extxyz_structures("full.extxyz", "train.extxyz", "val.extxyz", 0.1)
命令行参数详解
参数 |
缩写 |
描述 |
|---|---|---|
|
|
输入 EXTXYZ 文件路径（必需） |
|
|
训练集输出文件名（默认: train.extxyz） |
|
|
验证集输出文件名（默认: validation.extxyz） |
|
|
验证集比例（默认: 0.1） |
|
|
显示帮助信息 |
核心函数详解
split_extxyz_structures 函数
- split_extxyz_structures(input_file, train_file, validation_file, validation_ratio)
主功能函数：分割 EXTXYZ 结构数据集
- Parameters:
- Raises:
FileNotFoundError – 当输入文件不存在时
ValueError – 当验证集比例不在有效范围内时
- Returns:
None
函数执行流程：
文件读取与解析
structures = [] # 存储所有结构 current_struct = [] # 当前正在读取的结构 with open(input_file, 'r') as f: for line in f: stripped = line.strip() if stripped.isdigit(): # 检测结构开始（原子数行） if current_struct: # 保存前一个结构 structures.append(current_struct) current_struct = [] current_struct.append(line) else: # 结构内容行 current_struct.append(line) # 保存最后一个结构 if current_struct: structures.append(current_struct)
随机化处理
import random random.shuffle(structures) # 随机打乱结构顺序
比例计算与分割
\[ \begin{align}\begin{aligned}\begin{split}N_{\text{total}} = \text{len(structures)} \\ N_{\text{train}} = \lfloor N_{\text{total}} \times (1 - r_{\text{val}}) \rfloor \\ N_{\text{val}} = N_{\text{total}} - N_{\text{train}}\end{split}\\\begin{split}其中：\\ r_{\text{val}} = \text{validation_ratio}\end{split}\end{aligned}\end{align} \]split_idx = int(len(structures) * (1 - validation_ratio)) train_set = structures[:split_idx] # 训练集部分 val_set = structures[split_idx:] # 验证集部分
文件写入
# 写入训练集 with open(train_file, 'w') as f: for struct in train_set: f.write(''.join(struct)) # 合并结构行 # 写入验证集 with open(validation_file, 'w') as f: for struct in val_set: f.write(''.join(struct)) # 输出统计信息 print(f"训练集写入完成，样本数量: {len(train_set)}") print(f"验证集写入完成，样本数量: {len(val_set)}")
EXTXYZ 格式详解
格式概述
EXTXYZ（Extended XYZ）格式是标准 XYZ 格式的扩展，支持存储丰富的计算数据：
EXTXYZ 文件结构示例
# 结构1
2 # 原子数
Lattice="5.0 0.0 0.0 0.0 5.0 0.0 0.0 0.0 5.0" Properties=species:S:1:pos:R:3 energy=-10.5 free_energy=-10.5
W 0.000000 0.000000 0.000000 # 原子1
W 2.500000 2.500000 2.500000 # 原子2
# 结构2
3 # 原子数
Lattice="4.5 0.0 0.0 0.0 4.5 0.0 0.0 0.0 4.5" Properties=species:S:1:pos:R:3 energy=-15.2 free_energy=-15.2
W 0.000000 0.000000 0.000000
W 2.250000 2.250000 2.250000
W 0.000000 2.250000 2.250000
关键字段说明
字段 |
说明 |
|---|---|
原子数行 |
纯数字，表示该结构的原子总数 |
注释行 |
包含丰富的元数据： - Lattice：晶格向量（9个数字） - Properties：属性定义（物种、坐标等） - energy：总能（eV） - free_energy：自由能（eV） - 其他自定义属性 |
原子行 |
每行包含：元素符号 + X坐标 + Y坐标 + Z坐标 |
格式解析逻辑
模块使用以下逻辑识别结构边界：
结构开始：遇到纯数字行（原子数）
结构内容：直到下一个纯数字行之前的所有行
结构结束：遇到下一个结构的开始或文件结束
使用示例
基础示例
基础数据集分割
#!/usr/bin/env python3
"""基础数据集分割示例"""
from phymlp.tools.split_set import split_extxyz_structures
def main():
# 输入文件
input_file = "all_structures.extxyz"
# 输出文件
train_file = "training_set.extxyz"
validation_file = "validation_set.extxyz"
# 验证集比例（20%）
validation_ratio = 0.2
# 执行分割
split_extxyz_structures(
input_file=input_file,
train_file=train_file,
validation_file=validation_file,
validation_ratio=validation_ratio
)
print("数据集分割完成！")
print(f"训练集: {train_file}")
print(f"验证集: {validation_file}")
if __name__ == "__main__":
main()
高级使用示例
多比例分割实验
def experiment_with_different_ratios():
"""尝试不同的验证集比例"""
input_file = "dataset.extxyz"
ratios = [0.05, 0.1, 0.15, 0.2, 0.25]
for ratio in ratios:
# 生成文件名
train_file = f"train_{int(ratio*100)}pct.extxyz"
val_file = f"val_{int(ratio*100)}pct.extxyz"
print(f"\n使用验证集比例: {ratio*100}%")
# 执行分割
split_extxyz_structures(
input_file=input_file,
train_file=train_file,
validation_file=val_file,
validation_ratio=ratio
)
# 统计信息
count_structures(train_file, f"训练集 ({train_file})")
count_structures(val_file, f"验证集 ({val_file})")
def count_structures(filename, label):
"""统计EXTXYZ文件中的结构数"""
count = 0
with open(filename, 'r') as f:
for line in f:
if line.strip().isdigit():
count += 1
print(f" {label}: {count} 个结构")
集成到训练工作流
完整的数据准备流程
def prepare_training_data():
"""完整的训练数据准备流程"""
print("=" * 50)
print("数据准备流程开始")
print("=" * 50)
# 1. 收集原始数据
print("\n步骤1: 收集原始计算数据...")
raw_files = collect_vasp_results()
# 2. 转换为EXTXYZ格式
print("\n步骤2: 转换为EXTXYZ格式...")
for vasp_dir in raw_files:
convert_outcar_to_extxyz(vasp_dir)
# 3. 合并所有数据
print("\n步骤3: 合并数据集...")
merge_extxyz_files("raw_data/", "combined.extxyz")
# 4. 数据清洗
print("\n步骤4: 数据清洗...")
clean_dataset("combined.extxyz", "cleaned.extxyz")
# 5. 数据集分割
print("\n步骤5: 分割训练集和验证集...")
split_extxyz_structures(
input_file="cleaned.extxyz",
train_file="final_train.extxyz",
validation_file="final_validation.extxyz",
validation_ratio=0.15 # 15%验证集
)
# 6. 最终验证
print("\n步骤6: 验证数据集...")
validate_datasets("final_train.extxyz", "final_validation.extxyz")
print("\n" + "=" * 50)
print("数据准备完成！")
print("=" * 50)
带随机种子的分割
可重复的数据分割
import random
def split_with_random_seed(input_file, train_file, val_file, val_ratio, seed=42):
"""带随机种子的可重复分割"""
# 设置随机种子
random.seed(seed)
# 读取结构
structures = []
current_struct = []
with open(input_file, 'r') as f:
for line in f:
stripped = line.strip()
if stripped.isdigit():
if current_struct:
structures.append(current_struct)
current_struct = []
current_struct.append(line)
else:
current_struct.append(line)
if current_struct:
structures.append(current_struct)
# 打乱顺序（受随机种子控制）
random.shuffle(structures)
# 分割
split_idx = int(len(structures) * (1 - val_ratio))
train_set = structures[:split_idx]
val_set = structures[split_idx:]
# 写入文件
with open(train_file, 'w') as f:
for struct in train_set:
f.write(''.join(struct))
with open(val_file, 'w') as f:
for struct in val_set:
f.write(''.join(struct))
print(f"使用随机种子: {seed}")
print(f"训练集: {len(train_set)} 个结构")
print(f"验证集: {len(val_set)} 个结构")
return len(train_set), len(val_set)
错误处理机制
输入文件检查
import os
def validate_input_file(input_file):
"""验证输入文件"""
if not os.path.exists(input_file):
raise FileNotFoundError(f"输入文件不存在: {input_file}")
if not os.path.isfile(input_file):
raise ValueError(f"输入路径不是文件: {input_file}")
if os.path.getsize(input_file) == 0:
raise ValueError(f"输入文件为空: {input_file}")
比例参数验证
def validate_ratio(validation_ratio):
"""验证比例参数"""
if not isinstance(validation_ratio, (int, float)):
raise TypeError(f"验证集比例必须是数字，得到: {type(validation_ratio)}")
if validation_ratio < 0 or validation_ratio > 1:
raise ValueError(f"验证集比例必须在0和1之间，得到: {validation_ratio}")
if validation_ratio == 0 or validation_ratio == 1:
print(f"警告: 验证集比例为 {validation_ratio}，将导致一个集合为空")
格式验证
def validate_extxyz_format(filename):
"""验证EXTXYZ文件格式"""
structure_count = 0
line_count = 0
with open(filename, 'r') as f:
for line in f:
line_count += 1
stripped = line.strip()
if stripped.isdigit():
structure_count += 1
atom_count = int(stripped)
# 这里可以添加更多验证逻辑
print(f"文件包含 {structure_count} 个结构，共 {line_count} 行")
if structure_count == 0:
raise ValueError("文件不包含任何有效结构")
性能优化
内存优化策略
对于大型数据集，可以使用流式处理方法：
def stream_split_extxyz(input_file, train_file, val_file, val_ratio, seed=None):
"""流式处理大型EXTXYZ文件"""
if seed is not None:
random.seed(seed)
# 第一步：扫描文件，记录每个结构的起始位置和大小
structure_info = []
current_pos = 0
current_size = 0
with open(input_file, 'r') as f:
for line in f:
if line.strip().isdigit():
if current_size > 0:
structure_info.append((current_pos, current_size))
current_pos = f.tell() - len(line)
current_size = len(line)
else:
current_size += len(line)
if current_size > 0:
structure_info.append((current_pos, current_size))
# 第二步：随机选择验证集
total_structures = len(structure_info)
val_count = int(total_structures * val_ratio)
val_indices = random.sample(range(total_structures), val_count)
val_indices_set = set(val_indices)
# 第三步：流式写入
with open(input_file, 'r') as infile, \
open(train_file, 'w') as train_out, \
open(val_file, 'w') as val_out:
for i, (pos, size) in enumerate(structure_info):
infile.seek(pos)
structure_content = infile.read(size)
if i in val_indices_set:
val_out.write(structure_content)
else:
train_out.write(structure_content)
return total_structures - val_count, val_count
并行处理优化
from concurrent.futures import ProcessPoolExecutor
def parallel_structure_processing(structures, n_workers=4):
"""并行处理结构数据"""
def process_structure(struct):
# 处理单个结构的函数
# 这里可以添加能量计算、特征提取等
return processed_struct
with ProcessPoolExecutor(max_workers=n_workers) as executor:
processed_structures = list(executor.map(process_structure, structures))
return processed_structures
最佳实践建议
数据集比例选择
数据集大小 |
建议比例 |
说明 |
|---|---|---|
小型（< 1000） |
20%-30% |
需要更多验证数据 |
中型（1000-10000） |
15%-20% |
平衡训练和验证 |
大型（> 10000） |
5%-10% |
训练数据更重要 |
超大型（> 100000） |
1%-5% |
最小化验证损失 |
数据质量检查
分割前建议进行以下检查：
完整性检查：确保所有结构格式正确
能量范围：检查能量值是否在合理范围
力收敛：确认原子力已收敛
重复结构：移除高度相似的结构
异常检测：识别并处理异常值
随机性保证
为了确保可重复性：
固定随机种子：在科学研究中很重要
记录分割信息：保存分割日志
版本控制：数据集版本管理
常见问题解答
Q1: 分割时提示 “输入文件不存在”
解决方案：
# 检查文件路径
ls -la your_file.extxyz
# 使用绝对路径
python split_set.py --input /absolute/path/to/file.extxyz
Q2: 验证集比例设置错误
有效范围：0.0 到 1.0 之间的小数。
正确示例：
# 正确：10%验证集
phymlp-kit split_set --val_ratio 0.1
# 错误：百分比形式
phymlp-kit split_set --val_ratio 10 # 错误！
Q3: 分割后文件为空
可能原因：
输入文件为空
验证集比例为 0 或 1
文件格式不正确
检查方法：
def check_file_content(filename):
with open(filename, 'r') as f:
lines = f.readlines()
print(f"文件 {filename} 有 {len(lines)} 行")
# 显示前几行
for i, line in enumerate(lines[:5]):
print(f"行 {i+1}: {line.strip()}")
Q4: 需要按能量分层抽样
解决方案：修改分割函数，实现分层抽样：
def stratified_split_by_energy(input_file, train_file, val_file, val_ratio):
"""按能量分层抽样"""
# 1. 按能量分组结构
energy_groups = group_structures_by_energy(input_file)
train_structures = []
val_structures = []
# 2. 在每个能量组内按比例抽样
for energy, structures in energy_groups.items():
n_val = max(1, int(len(structures) * val_ratio))
val_indices = random.sample(range(len(structures)), n_val)
for i, struct in enumerate(structures):
if i in val_indices:
val_structures.append(struct)
else:
train_structures.append(struct)
# 3. 写入文件
write_structures(train_structures, train_file)
write_structures(val_structures, val_file)
API 参考
函数列表
- split_extxyz_structures(input_file, train_file, validation_file, validation_ratio)
（详见上文）
异常列表
- exception FileNotFoundError
当输入文件不存在时抛出。
- exception ValueError
当验证集比例不在有效范围内时抛出。
- exception TypeError
当参数类型不正确时抛出。
相关链接
quickstart - 快速开始指南
main_py - 主模块说明
get_KPATH 模块说明 - KPATH生成模块说明
更新日志
版本历史
版本 |
更新内容 |
|---|---|
v1.0.0 |
初始版本，基本分割功能 |
v1.1.0 |
添加命令行参数支持 |
v1.2.0 |
改进错误处理和格式验证 |
v1.3.0 |
添加性能优化选项 |
v1.4.0 |
支持随机种子设置 |
计划功能
智能分割：基于结构特征的智能分割策略
交叉验证：自动生成K折交叉验证数据集
数据增强：集成数据增强功能
可视化：数据集分布可视化工具
质量评估：数据集质量自动评估
Note
欢迎提交功能建议和改进意见！