LearnEP - NEP 机器学习势函数主动学习训练工具

LearnEP 是一个基于 NEP (Neuroevolution Potential) 的机器学习势函数主动学习训练工具。它集成了结构生成、FPS选择、MaxVol主动学习、NEP训练和可视化等完整功能，专为 HPC 环境下的科学计算工作流设计。

🚀 主要特性

🔬 完整的主动学习工作流：从初始数据准备到收敛的自动化流程
⚡ 高性能计算支持：PBS/SLURM 调度系统集成，GPU 加速计算
🎯 智能结构选择：FPS 最远点采样 + MaxVol 主动学习算法
📊 训练可视化：实时监控训练过程，生成专业图表
🔧 灵活配置：YAML 配置文件，支持迭代特定规则
🛡️ 容错机制：部分任务失败不终止整体流程，支持断点续传

📦 安装

环境要求

Python 3.8+
CUDA 支持的 GPU (推荐)
PBS 或 SLURM 调度系统 (HPC 环境)

使用 uv 安装 (推荐)

# 克隆仓库
git clone https://github.com/your-org/learnep.git
cd learnep

# 安装依赖
uv sync

# 激活环境
source .venv/bin/activate

使用 pip 安装

# 克隆仓库
git clone https://github.com/your-org/learnep.git
cd learnep

# 创建虚拟环境
python -m venv .venv
source .venv/bin/activate

# 安装依赖
pip install -e .

依赖项

核心依赖：numpy, jax, ase, typer, pyyaml
科学计算：scikit-learn, matplotlib
领域特定：pynep (NEP 势函数)
可视化：rich (CLI 界面)

🎯 快速开始

1. 准备配置文件

创建 config.yaml 配置文件：

global:
  work_dir: "./work"
  max_iterations: 50
  scheduler:
    type: "pbs"  # 或 "slurm"
    submit_cmd: "qsub {script}"
    status_cmd: "qstat {job_id}"
    cancel_cmd: "qdel {job_id}"

initial_data:
  # 方式1: 直接提供训练集
  train_data: "path/to/train.xyz"
  
  # 方式2: strgen + FPS + DFT
  strgen:
    enabled: true
    base_structure: "POSCAR"
    n_samples: 5000
    mode: "volumes"
  
  fps:
    enabled: true
    target_samples: 500
    min_distance: 0.05
    fallback_model: "nep89"
    fallback_model_path: "/path/to/nep89/nep.txt"

maxvol:
  gamma_min: 1.05
  gamma_max: 20.0
  max_select_per_round: 50
  device: "gpu"
  batch_size: 10000

nep:
  first_train: |
    type 2 Si O
    cutoff 6 5
    n_max 4 4
    basis_size 8 8
    l_max 4 2 1
    neuron 80
    batch 1000
    generation 500000
    population 50
    lambda_e 1.0
    lambda_f 1.0
    lambda_v 0.1

vasp:
  input_files:
    - "input/INCAR"
    - "input/POTCAR"
    - "input/KPOINTS"
  timeout: 86400
  max_parallel: 10

2. 运行主动学习工作流

# 运行完整的主动学习工作流
learnep run --config config.yaml

# 只运行 Phase 0 (初始数据准备)
learnep phase0 --config config.yaml

# 检查工作状态
learnep status --config config.yaml --detailed

# 从指定迭代恢复
learnep resume --config config.yaml --from-iter 5 --clean

# 验证配置文件
learnep check --config config.yaml

3. 使用独立工具

# FPS 最远点采样
learnep fps select --input structures.xyz --nep nep.txt --output selected.xyz

# MaxVol 主动学习
learnep maxvol select-gamma --nep nep.txt --input md.xyz --asi active_set.asi

# 结构生成
learnep strgen volumes --input POSCAR --output generated.xyz --n-samples 1000

# 训练可视化
learnep trainplot report --work-dir ./nep_training

📁 项目结构

learnep/
├── src/learnep/
│   ├── __init__.py          # 包初始化
│   ├── commands.py          # 主 CLI 入口
│   ├── config.yaml          # 默认配置
│   ├── plan.md             # 架构设计文档
│   ├── AGENTS.md           # 开发指南
│   │
│   ├── fps/                # 最远点采样
│   │   ├── __init__.py
│   │   ├── core.py         # FPS 算法实现
│   │   ├── commands.py     # CLI 接口
│   │   ├── visualize.py    # 可视化工具
│   │   └── README.md
│   │
│   ├── maxvol/             # MaxVol 主动学习
│   │   ├── __init__.py
│   │   ├── core.py         # JAX 加速算法
│   │   ├── commands.py     # CLI 接口
│   │   ├── tools.py        # 工具函数
│   │   ├── selector.py     # 选择逻辑
│   │   ├── io.py           # 文件 I/O
│   │   └── README.md
│   │
│   ├── strgen/             # 结构生成器
│   │   ├── __init__.py
│   │   ├── core.py         # 生成算法
│   │   ├── commands.py     # CLI 接口
│   │   └── README.md
│   │
│   ├── trainplot/          # 训练可视化
│   │   ├── __init__.py
│   │   ├── core.py         # 绘图核心功能
│   │   ├── commands.py     # CLI 接口
│   │   └── README.md
│   │
│   └── monitor/            # 监控和工作流管理
│       ├── __init__.py
│       ├── scheduler.py     # HPC 调度器接口
│       ├── monitor.py      # 核心监控器
│       ├── phase0.py       # Phase 0 执行器
│       ├── active_learning.py # 主动学习循环
│       └── README.md
│
├── docs/                    # 文档
├── tests/                   # 测试
├── examples/               # 示例
└── README.md              # 项目说明

🔧 核心模块

📊 TrainPlot - 训练可视化

自动解析 NEP 训练输出文件并生成专业图表：

# 生成完整训练报告
learnep trainplot report --work-dir ./nep_training

# 绘制特定图表
learnep trainplot loss --work-dir ./nep_training
learnep trainplot energy --work-dir ./nep_training --dataset train
learnep trainplot force --work-dir ./nep_training --dataset test

支持的文件格式：

loss.out - 损失函数和 RMSE 数据
energy_*.out - 能量预测 (eV/atom)
force_*.out - 力预测 (eV/Å)
virial_*.out - Virial 预测 (eV/atom)

🎯 FPS - 最远点采样

从大量结构中选择代表性子集：

learnep fps select \
  --input structures.xyz \
  --nep nep.txt \
  --output selected.xyz \
  --index selected_index.txt \
  --min-distance 0.05

🧠 MaxVol - 主动学习选择

基于 MaxVol 算法的智能结构选择：

# Gamma 扫描
learnep maxvol select-gamma \
  --nep nep.txt \
  --input md.xyz \
  --asi active_set.asi \
  --output high_gamma.xyz \
  --gamma-min 1.05

# 活跃集更新
learnep maxvol select-active \
  --train train.xyz \
  --nep nep.txt \
  --asi active_set.asi

# 扩展选择
learnep maxvol select-extend \
  --train train.xyz \
  --candidate high_gamma.xyz \
  --output extended.xyz

🏗️ StrGen - 结构生成器

多种结构生成算法：

# 体积变化
learnep strgen volumes \
  --input POSCAR \
  --output generated.xyz \
  --n-samples 1000 \
  --volume-min 0.9 \
  --volume-max 1.1

# 位移扰动
learnep strgen displacements \
  --input POSCAR \
  --output generated.xyz \
  --n-samples 1000 \
  --distance 0.01

# 标准生成
learnep strgen standard \
  --input POSCAR \
  --output generated.xyz \
  --n-samples 1000

🖥️ HPC 集成

PBS 调度系统

global:
  scheduler:
    type: "pbs"
    submit_cmd: "qsub {script}"
    status_cmd: "qstat {job_id}"
    cancel_cmd: "qdel {job_id}"
    check_interval: 60

SLURM 调度系统

global:
  scheduler:
    type: "slurm"
    submit_cmd: "sbatch {script}"
    status_cmd: "squeue -j {job_id}"
    cancel_cmd: "scancel {job_id}"
    check_interval: 60

作业脚本模板

系统支持命令注入，自动生成作业脚本：

nep:
  job_script: |
    #!/bin/bash
    #PBS -N nep_train_{iter}
    #PBS -q gpu
    #PBS -l nodes=1:ppn=1:gpus=1
    cd $PBS_O_WORKDIR
    nep > nep.log
    touch DONE

📈 工作流程示例

完整主动学习流程

# 1. 检查配置
learnep check --config config.yaml

# 2. 运行完整工作流
learnep run --config config.yaml

# 3. 监控进度
learnep status --config config.yaml --detailed

# 4. 如需恢复
learnep resume --config config.yaml --from-iter 10 --clean

分步执行

# Phase 0: 初始数据准备
learnep phase0 --config config.yaml

# 手动执行后续迭代
learnep workflow run --config config.yaml --from-iter 1

🛠️ 开发指南

代码风格

项目遵循严格的代码规范，详见 AGENTS.md：

导入顺序：标准库 → 第三方 → 本地模块
类型注解：所有函数参数和返回值
文档字符串：Google 风格，中文描述
错误处理：使用 logging，避免 print
命名规范：snake_case，描述性命名

测试

# 运行所有测试
pytest tests/

# 运行特定模块测试
pytest tests/test_fps.py

# 生成覆盖率报告
pytest --cov=learnep tests/

贡献指南

Fork 项目
创建功能分支
编写测试
提交 Pull Request
通过 CI 检查

📚 文档

架构设计 - 系统架构和设计理念
开发指南 - 代码规范和开发流程
模块文档 - 各模块详细说明
API 参考 - 完整的 API 文档

🤝 社区

问题报告：GitHub Issues
功能请求：GitHub Discussions
开发交流：Discord 频道

📄 许可证

本项目采用 MIT 许可证，详见 LICENSE 文件。

🙏 致谢

GPUMD 团队 - NEP 势函数和 GPUMD 软件
ASE 团队 - 原子模拟环境
JAX 团队 - 高性能计算框架
科学计算社区 - 算法和工具支持

LearnEP - 让 NEP 训练更智能、更高效！ 🚀

Name		Name	Last commit message	Last commit date
Latest commit History 98 Commits
src/learnep		src/learnep
.gitignore		.gitignore
AGENTS.md		AGENTS.md
README.md		README.md
pyproject.toml		pyproject.toml
uv.lock		uv.lock

gtiders/learnep

Folders and files

Latest commit

History

Repository files navigation