NEP框架重构
This commit is contained in:
121
src/workflow.py
121
src/workflow.py
@@ -50,75 +50,92 @@ class Workflow:
|
||||
if step_name == "00.md":
|
||||
step_dir = os.path.join(iter_path, "00.md")
|
||||
|
||||
# --- A. 第一轮初始化：POSCAR -> model.xyz ---
|
||||
# 1. 第一轮初始化：POSCAR -> model.xyz (保持不变)
|
||||
if iter_id == 0:
|
||||
os.makedirs(step_dir, exist_ok=True)
|
||||
poscar_name = self.param['files']['poscar']
|
||||
poscar_src = os.path.join(self.data_dir, poscar_name)
|
||||
|
||||
if os.path.exists(poscar_src):
|
||||
shutil.copy(poscar_src, os.path.join(step_dir, poscar_name))
|
||||
|
||||
atom_labels = self.param['files'].get('label', '')
|
||||
kit_path = self.machine.config['paths'].get('gpumdkit', 'gpumdkit.sh')
|
||||
|
||||
cmd = f"{kit_path} -addlabel {poscar_name} {atom_labels}"
|
||||
self.logger.info(f"Initializing model.xyz: {cmd}")
|
||||
subprocess.check_call(cmd, shell=True, cwd=step_dir)
|
||||
else:
|
||||
self.logger.error(f"POSCAR missing: {poscar_src}")
|
||||
continue # 跳过这一步
|
||||
|
||||
# 检查 model.xyz 是否存在 (无论是刚生成的，还是上一轮传过来的)
|
||||
# 注意：如果是 Iter > 0，这里需要逻辑从上一轮获取，目前先保证 Iter 0
|
||||
current_model_xyz = os.path.join(step_dir, "model.xyz")
|
||||
if not os.path.exists(current_model_xyz):
|
||||
self.logger.error(f"Critical: model.xyz not found in {step_dir}")
|
||||
continue
|
||||
|
||||
# --- B. 执行子任务 (Preheat / Production) ---
|
||||
sub_tasks = step_conf.get('sub_tasks', [])
|
||||
if not sub_tasks:
|
||||
self.logger.warning("No sub_tasks defined for 00.md in param.yaml!")
|
||||
|
||||
for sub in sub_tasks:
|
||||
template_sub_name = sub.get('template_sub')
|
||||
if not template_sub_name: continue
|
||||
|
||||
# 1. 定义子目录 (e.g., iter_00/00.md/preheat)
|
||||
sub_work_dir = os.path.join(step_dir, template_sub_name)
|
||||
os.makedirs(sub_work_dir, exist_ok=True)
|
||||
|
||||
self.logger.info(f"--> Setup Sub-task: {template_sub_name}")
|
||||
|
||||
# 2. 【关键】将 model.xyz 分发到子目录
|
||||
shutil.copy(current_model_xyz, os.path.join(sub_work_dir, "model.xyz"))
|
||||
|
||||
# 3. 准备 Template (run.in)
|
||||
# 路径: template/00.md/preheat/
|
||||
template_path = os.path.join(self.template_dir, "00.md", template_sub_name)
|
||||
if not os.path.exists(template_path):
|
||||
self.logger.error(f"Template not found: {template_path}")
|
||||
continue
|
||||
|
||||
# 4. 实例化 MD 任务
|
||||
# 这里的 name 仅仅是日志用的标识
|
||||
md_task = MDStep(f"MD-{template_sub_name}", sub_work_dir, self.machine, self.config)
|
||||
# 确保 gpumdkit 路径可用
|
||||
kit_path = self.machine.config['paths'].get('gpumdkit', 'gpumdkit.sh')
|
||||
|
||||
# 5. 运行 (传入势函数路径 和 模板路径)
|
||||
# run() 方法内部会处理 nep.txt 的复制和 run.in 的复制
|
||||
success = md_task.run(self.current_nep_pot, template_path)
|
||||
# ----------------------------------------------------
|
||||
# 2. 核心修改：分别处理 preheat 和 production
|
||||
# ----------------------------------------------------
|
||||
|
||||
if success:
|
||||
# 记录最新的 dump.xyz 位置，供下一步 Select 使用
|
||||
self.last_dump_path = os.path.join(sub_work_dir, "dump.xyz")
|
||||
# 同时也把最新的 model.xyz (如果MD改变了结构) 或 dump 更新为下一次的起点？
|
||||
# 目前逻辑：Select 步会用 self.last_dump_path
|
||||
else:
|
||||
self.logger.error(f"MD Sub-task {template_sub_name} Failed.")
|
||||
# 根据策略决定是否中断，这里暂时中断
|
||||
break
|
||||
# ==========================
|
||||
# === Sub-task 1: Preheat (预热) ===
|
||||
# 逻辑：复制model.xyz -> 跑MD -> 跑201采样 -> 生成 sampled_structures.xyz
|
||||
preheat_dir = os.path.join(step_dir, "preheat")
|
||||
os.makedirs(preheat_dir, exist_ok=True)
|
||||
|
||||
# 准备文件
|
||||
current_model_xyz = os.path.join(step_dir, "model.xyz")
|
||||
shutil.copy(current_model_xyz, os.path.join(preheat_dir, "model.xyz"))
|
||||
shutil.copy(self.current_nep_pot, os.path.join(preheat_dir, "nep.txt"))
|
||||
shutil.copy(os.path.join(self.template_dir, "00.md", "preheat", "run.in"),
|
||||
os.path.join(preheat_dir, "run.in"))
|
||||
|
||||
self.logger.info(">>> Running Preheat MD...")
|
||||
# 使用 Machine 运行 GPUMD (假设 machine.yaml 里 gpumd 是基础命令)
|
||||
self.machine.execute("gpumd", preheat_dir)
|
||||
|
||||
# [关键] Preheat 后处理：采样
|
||||
if os.path.exists(os.path.join(preheat_dir, "dump.xyz")):
|
||||
self.logger.info(">>> Running Sampling (201)...")
|
||||
# 构造命令: echo -e "201\ndump.xyz\nuniform\n4" | gpumdkit.sh
|
||||
# 注意：根据你的描述 "dump.xyz uniform 4"，我这里构造输入流
|
||||
# 如果你的脚本交互顺序不同，请调整这里的字符串
|
||||
# 这里的 \n 代表回车
|
||||
input_str = "201\ndump.xyz\nuniform\n4"
|
||||
|
||||
try:
|
||||
# 调用 gpumdkit
|
||||
process = subprocess.Popen(
|
||||
kit_path,
|
||||
shell=True,
|
||||
cwd=preheat_dir,
|
||||
stdin=subprocess.PIPE,
|
||||
stdout=subprocess.PIPE,
|
||||
stderr=subprocess.PIPE,
|
||||
text=True
|
||||
)
|
||||
stdout, stderr = process.communicate(input=input_str)
|
||||
|
||||
if os.path.exists(os.path.join(preheat_dir, "sampled_structures.xyz")):
|
||||
self.logger.info("Sampled structures generated successfully.")
|
||||
else:
|
||||
self.logger.error(f"Sampling failed. Output: {stdout} Error: {stderr}")
|
||||
continue # 如果没生成采样文件，后续Production没法做
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error executing sampling: {e}")
|
||||
continue
|
||||
else:
|
||||
self.logger.error("Preheat dump.xyz missing.")
|
||||
continue
|
||||
|
||||
# === Sub-task 2: Production (加工/正式采样) ===
|
||||
# 逻辑：链接 sampled_structures -> 跑302 -> 跑presub.sh
|
||||
prod_dir = os.path.join(step_dir, "production")
|
||||
os.makedirs(prod_dir, exist_ok=True)
|
||||
|
||||
# 1. 建立软链接 (sampled_structures.xyz)
|
||||
src_sample = os.path.abspath(os.path.join(preheat_dir, "sampled_structures.xyz"))
|
||||
dst_sample = os.path.join(prod_dir, "sampled_structures.xyz")
|
||||
if os.path.exists(dst_sample): os.remove(dst_sample) # 清理旧的
|
||||
os.symlink(src_sample, dst_sample)
|
||||
|
||||
# 2. 准备基础文件
|
||||
self.logger.error(f"presub.sh execution failed: {e}")
|
||||
# Step: 01.select
|
||||
# ==========================
|
||||
elif step_name == "01.select":
|
||||
|
||||
Reference in New Issue
Block a user