564 lines
32 KiB
Python
564 lines
32 KiB
Python
# src/workflow.py
|
||
import os
|
||
import shutil
|
||
import logging
|
||
import subprocess
|
||
|
||
from src.utils import load_yaml, run_cmd_with_log
|
||
from src.machine import MachineManager
|
||
from src.state import StateTracker # 新增
|
||
from src.steps import MDStep, SelectStep, SCFStep, TrainStep
|
||
|
||
|
||
class Workflow:
|
||
def __init__(self, root_dir):
|
||
self.root_dir = root_dir
|
||
self.param = load_yaml(os.path.join(root_dir, "config/param.yaml"))
|
||
self.machine = MachineManager(os.path.join(root_dir, "config/machine.yaml"))
|
||
|
||
self.workspace = os.path.join(root_dir, "workspace")
|
||
self.data_dir = os.path.join(root_dir, "data")
|
||
self.template_dir = os.path.join(root_dir, "template")
|
||
|
||
self.logger = logging.getLogger()
|
||
|
||
# 初始化状态追踪
|
||
os.makedirs(self.workspace, exist_ok=True)
|
||
self.tracker = StateTracker(self.workspace)
|
||
|
||
# 初始变量
|
||
self.current_nep_pot = os.path.join(self.data_dir, self.param['files']['initial_pot'])
|
||
self.current_train_set = os.path.join(self.workspace, "accumulated_train.xyz")
|
||
|
||
def run(self):
|
||
self.logger.info(f"Workflow Started: {self.param['project']}")
|
||
|
||
for iteration in self.param['iterations']:
|
||
iter_id = iteration['id']
|
||
iter_name = f"iter_{iter_id:02d}"
|
||
iter_path = os.path.join(self.workspace, iter_name)
|
||
self.logger.info(f"\n >>> Processing Iteration: {iter_id} <<<")
|
||
os.makedirs(iter_path, exist_ok=True)
|
||
|
||
for step_conf in iteration['steps']:
|
||
step_name = step_conf['name']
|
||
|
||
# ==========================
|
||
# Step: 00.md
|
||
# ==========================
|
||
if step_name == "00.md":
|
||
step_dir = os.path.join(iter_path, "00.md")
|
||
|
||
# 1. 初始化 model.xyz (仅做一次)
|
||
task_id_init = f"{iter_name}.00.md.init"
|
||
|
||
if iter_id == 0:
|
||
if not self.tracker.is_done(task_id_init):
|
||
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...")
|
||
|
||
if run_cmd_with_log(cmd, step_dir, "init.log"):
|
||
self.tracker.mark_done(task_id_init)
|
||
else:
|
||
self.logger.error("Initialization failed. Check iter_00/00.md/init.log")
|
||
return
|
||
else:
|
||
self.logger.error("POSCAR missing.")
|
||
return
|
||
else:
|
||
self.logger.info("Skipping Init (Already Done).")
|
||
else:
|
||
# --- [新增逻辑] 后续轮次：从上一轮复制 model.xyz ---
|
||
|
||
# 1. 【核心修复】必须先确保目标文件夹存在！
|
||
if not os.path.exists(step_dir):
|
||
os.makedirs(step_dir, exist_ok=True)
|
||
|
||
# 2. 检查当前目录是否已有 model.xyz
|
||
if not os.path.exists(os.path.join(step_dir, "model.xyz")):
|
||
prev_iter_name = f"iter_{iter_id - 1:02d}"
|
||
prev_model_src = os.path.join(self.workspace, prev_iter_name, "00.md", "model.xyz")
|
||
|
||
if os.path.exists(prev_model_src):
|
||
self.logger.info(f"Copying model.xyz from {prev_iter_name}...")
|
||
shutil.copy(prev_model_src, os.path.join(step_dir, "model.xyz"))
|
||
else:
|
||
self.logger.error(f"Previous model.xyz not found: {prev_model_src}")
|
||
return
|
||
# 确保 gpumdkit 路径可用
|
||
kit_path = self.machine.config['paths'].get('gpumdkit', 'gpumdkit.sh')
|
||
|
||
# === Sub-task 1: Preheat ===
|
||
task_id_preheat = f"{iter_name}.00.md.preheat"
|
||
preheat_dir = os.path.join(step_dir, "preheat")
|
||
|
||
if not self.tracker.is_done(task_id_preheat):
|
||
self.logger.info(">>> Starting Preheat...")
|
||
os.makedirs(preheat_dir, exist_ok=True)
|
||
|
||
# 准备文件
|
||
shutil.copy(os.path.join(step_dir, "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"))
|
||
|
||
# A. 运行 GPUMD
|
||
# 假设 gpumd 命令直接运行，无输入
|
||
if not run_cmd_with_log("gpumd", preheat_dir, "step_exec.log"):
|
||
self.logger.error("Preheat GPUMD failed.")
|
||
return
|
||
|
||
# B. 运行 采样 (201)
|
||
# [修正] 严格按照要求: "201\ndump.xyz uniform 4" (中间无额外换行)
|
||
input_str_201 = "201\ndump.xyz uniform 4"
|
||
self.logger.info(">>> Running Sampling (201)...")
|
||
|
||
if run_cmd_with_log(kit_path, preheat_dir, "step_exec.log", input_str=input_str_201):
|
||
if os.path.exists(os.path.join(preheat_dir, "sampled_structures.xyz")):
|
||
self.tracker.mark_done(task_id_preheat)
|
||
else:
|
||
self.logger.error("sampled_structures.xyz not generated.")
|
||
return
|
||
else:
|
||
self.logger.error("Sampling command failed.")
|
||
return
|
||
else:
|
||
self.logger.info("Skipping Preheat (Already Done).")
|
||
|
||
# === Sub-task 2: Production ===
|
||
task_id_prod = f"{iter_name}.00.md.production"
|
||
prod_dir = os.path.join(step_dir, "production")
|
||
|
||
if not self.tracker.is_done(task_id_prod):
|
||
self.logger.info(">>> Starting Production...")
|
||
os.makedirs(prod_dir, exist_ok=True)
|
||
|
||
# 1. 准备基础文件到 production 根目录
|
||
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)
|
||
|
||
shutil.copy(self.current_nep_pot, os.path.join(prod_dir, "nep.txt"))
|
||
shutil.copy(os.path.join(self.template_dir, "00.md", "production", "run.in"),
|
||
os.path.join(prod_dir, "run.in"))
|
||
|
||
# 2. 运行 302 (生成 md 文件夹, sample_* 文件夹, presub.sh)
|
||
# 302 通常会读取当前目录的 run.in 并在 md/ 下生成拆分后的 run_x.in
|
||
input_str_302 = "302"
|
||
if not run_cmd_with_log(kit_path, prod_dir, "step_exec.log", input_str=input_str_302):
|
||
self.logger.error("302 command failed.")
|
||
return
|
||
|
||
if not os.path.exists(os.path.join(prod_dir, "presub.sh")):
|
||
self.logger.error("presub.sh not found.")
|
||
return
|
||
|
||
# ---------------------------------------------------------
|
||
# [新增] 3. 补全文件：将 nep.txt 和 run.in 复制到 md 文件夹
|
||
# ---------------------------------------------------------
|
||
md_subdir = os.path.join(prod_dir, "md")
|
||
if os.path.exists(md_subdir):
|
||
self.logger.info("Copying nep.txt and run.in to 'md' folder...")
|
||
shutil.copy(os.path.join(prod_dir, "nep.txt"), os.path.join(md_subdir, "nep.txt"))
|
||
# 复制 run.in，虽然 302 可能已经生成了 run_1.in 等，但为了保险或用户习惯，我们也拷进去
|
||
shutil.copy(os.path.join(prod_dir, "run.in"), os.path.join(md_subdir, "run.in"))
|
||
shutil.copy(os.path.join(prod_dir, "run.in"), os.path.join(md_subdir, "run_1.in"))
|
||
shutil.copy(os.path.join(prod_dir, "run.in"), os.path.join(md_subdir, "run_2.in"))
|
||
shutil.copy(os.path.join(prod_dir, "run.in"), os.path.join(md_subdir, "run_3.in"))
|
||
shutil.copy(os.path.join(prod_dir, "run.in"), os.path.join(md_subdir, "run_4.in"))
|
||
else:
|
||
self.logger.error("'md' folder was not created by 302 command.")
|
||
return
|
||
|
||
# 4. 运行 presub.sh
|
||
os.chmod(os.path.join(prod_dir, "presub.sh"), 0o755)
|
||
self.logger.info(">>> Executing presub.sh...")
|
||
|
||
if not run_cmd_with_log("./presub.sh", prod_dir, "step_exec.log"):
|
||
self.logger.error("presub.sh execution failed.")
|
||
return
|
||
|
||
# 5. 合并 dump
|
||
self.logger.info("Merging dump files...")
|
||
run_cmd_with_log("cat sample_*/dump.xyz > dump.xyz", prod_dir, "step_exec.log")
|
||
|
||
self.last_dump_path = os.path.join(prod_dir, "dump.xyz")
|
||
self.tracker.mark_done(task_id_prod)
|
||
else:
|
||
self.logger.info("Skipping Production (Already Done).")
|
||
self.last_dump_path = os.path.join(prod_dir, "dump.xyz")
|
||
|
||
# ==========================
|
||
# Step: 01.select
|
||
# ==========================
|
||
elif step_name == "01.select":
|
||
step_dir = os.path.join(iter_path, "01.select")
|
||
task_id_select = f"{iter_name}.01.select"
|
||
|
||
if not self.tracker.is_done(task_id_select):
|
||
self.logger.info(f"=== Step: 01.select ({step_conf.get('method', 'distance')}) ===")
|
||
os.makedirs(step_dir, exist_ok=True)
|
||
|
||
# 1. 准备 dump.xyz (软链接)
|
||
# 优先使用上一这步记录的路径，如果没有则尝试推断
|
||
dump_src = getattr(self, 'last_dump_path', None)
|
||
if not dump_src:
|
||
# 尝试推断：iter_XX/00.md/production/dump.xyz
|
||
dump_src = os.path.join(iter_path, "00.md", "production", "dump.xyz")
|
||
|
||
if os.path.exists(dump_src):
|
||
dst_dump = os.path.join(step_dir, "dump.xyz")
|
||
if os.path.exists(dst_dump): os.remove(dst_dump)
|
||
# 使用绝对路径建立软链
|
||
os.symlink(os.path.abspath(dump_src), dst_dump)
|
||
else:
|
||
self.logger.error(f"Source dump.xyz not found: {dump_src}")
|
||
return
|
||
|
||
# 2. 准备 nep.txt
|
||
shutil.copy(self.current_nep_pot, os.path.join(step_dir, "nep.txt"))
|
||
|
||
# 3. 准备 train.xyz
|
||
# 逻辑：第一轮用 model.xyz，之后用 self.current_train_set
|
||
if iter_id == 0:
|
||
# 第一轮：找同轮次 00.md 下的 model.xyz
|
||
model_xyz_src = os.path.join(iter_path, "00.md", "model.xyz")
|
||
if os.path.exists(model_xyz_src):
|
||
shutil.copy(model_xyz_src, os.path.join(step_dir, "train.xyz"))
|
||
else:
|
||
self.logger.error("model.xyz not found for initial train.xyz")
|
||
return
|
||
else:
|
||
# 后续轮次
|
||
if os.path.exists(self.current_train_set):
|
||
shutil.copy(self.current_train_set, os.path.join(step_dir, "train.xyz"))
|
||
else:
|
||
self.logger.error(f"Previous train set missing: {self.current_train_set}")
|
||
return
|
||
|
||
# 4. 执行筛选逻辑
|
||
method = step_conf.get('method', 'distance')
|
||
params = step_conf.get('params', [0.01, 60, 120]) # [threshold, min, max]
|
||
kit_path = self.machine.config['paths'].get('gpumdkit', 'gpumdkit.sh')
|
||
|
||
# === 分支 A: 按距离筛选 (Loop) ===
|
||
if method == "distance":
|
||
threshold = params[0]
|
||
target_min, target_max = params[1], params[2]
|
||
step_size = 0.001
|
||
max_attempts = 15
|
||
success = False
|
||
|
||
self.logger.info(f"Targeting {target_min}-{target_max} structures. Initial thr={threshold}")
|
||
|
||
for attempt in range(max_attempts):
|
||
self.logger.info(f"--- Attempt {attempt + 1}: Threshold {threshold:.4f} ---")
|
||
|
||
# 构造输入: 203 -> filenames -> 1 (distance) -> threshold
|
||
# 格式: "203\ndump.xyz train.xyz nep.txt\n1\n0.01"
|
||
input_str = f"203\ndump.xyz train.xyz nep.txt\n1\n{threshold}"
|
||
|
||
# 运行 gpumdkit (输出重定向到 append 模式的日志)
|
||
if not run_cmd_with_log(kit_path, step_dir, "select_exec.log", input_str=input_str):
|
||
self.logger.error("gpumdkit execution failed.")
|
||
break # 致命错误直接退出
|
||
|
||
# 检查 selected.xyz 数量
|
||
if not os.path.exists(os.path.join(step_dir, "selected.xyz")):
|
||
self.logger.warning("selected.xyz not found. Retrying...")
|
||
continue
|
||
|
||
try:
|
||
# grep -c "Lat" selected.xyz
|
||
count_out = subprocess.check_output('grep -c "Lat" selected.xyz', shell=True,
|
||
cwd=step_dir)
|
||
count = int(count_out.decode().strip())
|
||
except:
|
||
count = 0
|
||
|
||
self.logger.info(f"Selected count: {count}")
|
||
|
||
# 判断逻辑
|
||
if target_min <= count <= target_max:
|
||
self.logger.info(f"Success! {count} structures selected.")
|
||
success = True
|
||
break
|
||
elif count < target_min:
|
||
# 选少了 -> 阈值太严 -> 降低阈值? (通常距离阈值越低，选的越多？)
|
||
# 假设：Threshold 代表“最小允许距离”。距离 > Thr 才选。
|
||
# 那么 Thr 越小，条件越松，选的越多。
|
||
self.logger.info("Too few. Decreasing threshold (Loosening).")
|
||
threshold -= step_size
|
||
if threshold < 0: threshold = 0.0001
|
||
else: # count > target_max
|
||
# 选多了 -> 提高阈值
|
||
self.logger.info("Too many. Increasing threshold (Tightening).")
|
||
threshold += step_size
|
||
|
||
if success:
|
||
self.tracker.mark_done(task_id_select)
|
||
else:
|
||
self.logger.warning(
|
||
"Failed to reach target count within max attempts. Using last result.")
|
||
# 这里可以选择是否标记完成，或者报错暂停。目前标记完成以便继续。
|
||
self.tracker.mark_done(task_id_select)
|
||
|
||
# === 分支 B: 按个数/随机筛选 (One-shot) ===
|
||
elif method == "random":
|
||
min_n, max_n = params[1], params[2]
|
||
# 构造输入: 203 -> filenames -> 2 (number) -> min max
|
||
input_str = f"203\ndump.xyz train.xyz nep.txt\n2\n{min_n} {max_n}"
|
||
|
||
self.logger.info(f"Random selection: {min_n}-{max_n}")
|
||
if run_cmd_with_log(kit_path, step_dir, "select_exec.log", input_str=input_str):
|
||
self.tracker.mark_done(task_id_select)
|
||
else:
|
||
self.logger.error("Random selection failed.")
|
||
return
|
||
|
||
else:
|
||
self.logger.info("Skipping Select (Already Done).")
|
||
# ==========================
|
||
# Step: 02.scf (VASP Calculation)
|
||
# ==========================
|
||
elif step_name == "02.scf":
|
||
step_dir = os.path.join(iter_path, "02.scf")
|
||
task_id_scf = f"{iter_name}.02.scf"
|
||
|
||
if not self.tracker.is_done(task_id_scf):
|
||
self.logger.info("=== Step: 02.scf (VASP) ===")
|
||
os.makedirs(step_dir, exist_ok=True)
|
||
|
||
# 1. 准备 selected.xyz
|
||
# 尝试从同轮次的 01.select 获取
|
||
select_step_dir = os.path.join(iter_path, "01.select")
|
||
src_selected = os.path.join(select_step_dir, "selected.xyz")
|
||
if not os.path.exists(src_selected):
|
||
self.logger.error(f"selected.xyz not found in {select_step_dir}")
|
||
return
|
||
|
||
dst_selected = os.path.join(step_dir, "selected.xyz")
|
||
if os.path.exists(dst_selected): os.remove(dst_selected)
|
||
os.symlink(os.path.abspath(src_selected), dst_selected)
|
||
|
||
# 2. 运行 301 拆分结构
|
||
# 命令: echo -e "301\niter" | gpumdkit.sh
|
||
# 这会生成 iterX_1, iterX_2... 和 fp 文件夹
|
||
kit_path = self.machine.config['paths'].get('gpumdkit', 'gpumdkit.sh')
|
||
input_str_301 = "301\niter" # 这里 "iter" 是文件夹前缀名，gpumdkit 会自动加数字
|
||
|
||
self.logger.info("Splitting structures (301)...")
|
||
if not run_cmd_with_log(kit_path, step_dir, "scf_setup.log", input_str=input_str_301):
|
||
self.logger.error("301 command failed.")
|
||
return
|
||
|
||
# 3. 准备 VASP 输入文件到 'fp' 文件夹
|
||
# gpumdkit 生成的 fp 文件夹通常存放公共文件，子文件夹会软链过去
|
||
fp_dir = os.path.join(step_dir, "fp")
|
||
if not os.path.exists(fp_dir):
|
||
self.logger.error("'fp' directory was not created by 301.")
|
||
return
|
||
|
||
self.logger.info("Distributing VASP inputs to 'fp' folder...")
|
||
|
||
# A. POTCAR (来自 Data)
|
||
potcar_src = os.path.join(self.data_dir, self.param['files']['potcar'])
|
||
if os.path.exists(potcar_src):
|
||
shutil.copy(potcar_src, os.path.join(fp_dir, "POTCAR"))
|
||
else:
|
||
self.logger.error(f"POTCAR missing: {potcar_src}")
|
||
return
|
||
|
||
# B. INCAR (来自 Template)
|
||
# Template 路径: template/02.scf/INCAR
|
||
incar_src = os.path.join(self.template_dir, "02.scf", "INCAR")
|
||
if os.path.exists(incar_src):
|
||
shutil.copy(incar_src, os.path.join(fp_dir, "INCAR"))
|
||
else:
|
||
self.logger.error(f"INCAR missing in template: {incar_src}")
|
||
return
|
||
|
||
# C. KPOINTS (来自 Template, 可选)
|
||
kpoints_src = os.path.join(self.template_dir, "02.scf", "KPOINTS")
|
||
if os.path.exists(kpoints_src):
|
||
shutil.copy(kpoints_src, os.path.join(fp_dir, "KPOINTS"))
|
||
else:
|
||
self.logger.info("KPOINTS not found in template, assuming KSPACING in INCAR.")
|
||
|
||
# 4. 生成并提交计算任务
|
||
# 这里我们不理会 gpumdkit 生成的 presub.sh，而是根据 machine.yaml 生成自己的
|
||
executor_name = step_conf.get('executor', 'vasp_gpu') # 默认用 cpu
|
||
|
||
# 获取执行命令 (例如 "mpirun -np 32 vasp_std")
|
||
# 这里的逻辑需要调用 machine 模块的一个新功能：批量生成提交脚本
|
||
# 但为了简化，我们在 Local 模式下生成一个遍历脚本
|
||
|
||
self.logger.info(f"Generating batch submission script for {executor_name}...")
|
||
|
||
# 读取 machine 配置里的命令
|
||
exec_conf = self.machine.config['executors'].get(executor_name, {})
|
||
vasp_cmd = exec_conf.get('cmd', 'mpirun -np 1 vasp_std') # 默认值
|
||
|
||
# 生成 run_vasp.sh
|
||
run_script_path = os.path.join(step_dir, "run_vasp.sh")
|
||
with open(run_script_path, 'w') as f:
|
||
f.write("#!/bin/bash\n")
|
||
# 遍历 iter* 目录
|
||
f.write(f"for dir in iter*_*; do\n")
|
||
f.write(f" if [ -d \"$dir\" ]; then\n")
|
||
f.write(f" echo \"Running VASP in $dir ...\"\n")
|
||
f.write(f" cd $dir\n")
|
||
# 写入具体的 VASP 执行命令
|
||
f.write(f" {vasp_cmd} > vasp.log 2>&1\n") # 重定向日志
|
||
f.write(f" cd ..\n")
|
||
f.write(f" fi\n")
|
||
f.write(f"done\n")
|
||
|
||
os.chmod(run_script_path, 0o755)
|
||
|
||
# 执行 VASP 计算
|
||
# 注意：如果是在 Slurm 上，这里应该提交 run_vasp.sh，并使用 Job ID 等待
|
||
# 目前 Local 模式直接运行
|
||
self.logger.info(">>> Executing VASP batch calculations (this may take time)...")
|
||
if not run_cmd_with_log("./run_vasp.sh", step_dir, "scf_exec.log"):
|
||
self.logger.error("VASP batch execution failed.")
|
||
return
|
||
|
||
# 5. 结果收集 (out2xyz)
|
||
self.logger.info("Collecting results (out2xyz)...")
|
||
cmd_collect = f"{kit_path} -out2xyz ."
|
||
if run_cmd_with_log(cmd_collect, step_dir, "scf_collect.log"):
|
||
# 检查结果
|
||
res_dir = os.path.join(step_dir, "NEPdataset-multiple_frames")
|
||
res_file = os.path.join(res_dir, "NEP-dataset.xyz")
|
||
|
||
if os.path.exists(res_file):
|
||
self.logger.info(f"VASP data collected: {res_file}")
|
||
# 保存这个路径供 Train 使用
|
||
self.new_data_chunk = res_file
|
||
self.tracker.mark_done(task_id_scf)
|
||
else:
|
||
self.logger.error("NEP-dataset.xyz not found after collection.")
|
||
else:
|
||
self.logger.error("out2xyz failed.")
|
||
|
||
else:
|
||
self.logger.info("Skipping SCF (Already Done).")
|
||
# 即使跳过，也要尝试恢复 self.new_data_chunk 变量，防止 Train 找不到数据
|
||
# 这里简单推断一下
|
||
res_file = os.path.join(step_dir, "NEPdataset-multiple_frames", "NEP-dataset.xyz")
|
||
if os.path.exists(res_file):
|
||
self.new_data_chunk = res_file
|
||
# ==========================
|
||
# Step: 03.train (Training)
|
||
# ==========================
|
||
elif step_name == "03.train":
|
||
step_dir = os.path.join(iter_path, "03.train")
|
||
task_id_train = f"{iter_name}.03.train"
|
||
|
||
if not self.tracker.is_done(task_id_train):
|
||
self.logger.info("=== Step: 03.train (NEP) ===")
|
||
os.makedirs(step_dir, exist_ok=True)
|
||
|
||
# 1. 准备 train.xyz (合并数据)
|
||
# 逻辑：Current Total Train = Previous Total Train + New Data Chunk
|
||
self.logger.info("Merging training data...")
|
||
|
||
train_xyz_path = os.path.join(step_dir, "train.xyz")
|
||
|
||
# 打开目标文件准备写入
|
||
with open(train_xyz_path, 'w') as outfile:
|
||
# A. 写入旧数据 (如果存在)
|
||
if os.path.exists(self.current_train_set):
|
||
self.logger.info(f"Appending previous data: {self.current_train_set}")
|
||
with open(self.current_train_set, 'r') as infile:
|
||
shutil.copyfileobj(infile, outfile)
|
||
|
||
# B. 写入新数据 (来自本轮 SCF)
|
||
# 尝试获取变量，如果变量丢失则尝试从路径推断
|
||
new_data = getattr(self, 'new_data_chunk', None)
|
||
if not new_data:
|
||
# 推断路径: iter_XX/02.scf/NEPdataset-multiple_frames/NEP-dataset.xyz
|
||
new_data = os.path.join(iter_path, "02.scf", "NEPdataset-multiple_frames",
|
||
"NEP-dataset.xyz")
|
||
|
||
if new_data and os.path.exists(new_data):
|
||
self.logger.info(f"Appending new data: {new_data}")
|
||
with open(new_data, 'r') as infile:
|
||
shutil.copyfileobj(infile, outfile)
|
||
else:
|
||
if iter_id == 0 and not os.path.exists(self.current_train_set):
|
||
self.logger.error("No training data available (neither previous nor new).")
|
||
return
|
||
else:
|
||
self.logger.warning("No new data found from SCF step. Training on old data only.")
|
||
|
||
# 更新全局变量指向最新的 train.xyz，供下一轮使用
|
||
self.current_train_set = train_xyz_path
|
||
|
||
# 2. 准备 nep.in
|
||
template_nep_in = os.path.join(self.template_dir, "03.train", "nep.in")
|
||
if os.path.exists(template_nep_in):
|
||
shutil.copy(template_nep_in, os.path.join(step_dir, "nep.in"))
|
||
else:
|
||
self.logger.error(f"nep.in template missing: {template_nep_in}")
|
||
return
|
||
|
||
# 3. 执行训练
|
||
executor_name = step_conf.get('executor', 'nep_local')
|
||
# 获取nep命令，比如 "nep" 或者 "/path/to/nep"
|
||
# 注意：nep 命令通常不需要参数，它会自动读取 nep.in
|
||
|
||
self.logger.info(f"Starting NEP training using {executor_name}...")
|
||
# 这里的 log 文件叫 train_exec.log
|
||
if not run_cmd_with_log("nep", step_dir,
|
||
"train_exec.log"): # 假设 cmd 是 nep，如果 machine.yaml 里有特殊定义请调整
|
||
self.logger.error("NEP training failed.")
|
||
return
|
||
|
||
# 检查是否生成了 nep.txt
|
||
if os.path.exists(os.path.join(step_dir, "nep.txt")):
|
||
self.logger.info("Training finished. nep.txt generated.")
|
||
# 更新全局势函数路径，供下一轮 MD 使用
|
||
self.current_nep_pot = os.path.join(step_dir, "nep.txt")
|
||
else:
|
||
self.logger.error("nep.txt not found after training.")
|
||
return
|
||
|
||
# 4. 后处理：绘图与归档
|
||
self.logger.info("Generating plots (gpumdkit.sh -plt train)...")
|
||
kit_path = self.machine.config['paths'].get('gpumdkit', 'gpumdkit.sh')
|
||
cmd_plt = f"{kit_path} -plt train"
|
||
|
||
run_cmd_with_log(cmd_plt, step_dir, "plot.log")
|
||
|
||
# [修改] 创建 output 目录在当前 iter 内部
|
||
output_dir = os.path.join(iter_path, "05.output")
|
||
os.makedirs(output_dir, exist_ok=True)
|
||
|
||
# 移动 png 图片
|
||
for file in os.listdir(step_dir):
|
||
if file.endswith(".png"):
|
||
src_png = os.path.join(step_dir, file)
|
||
dst_png = os.path.join(output_dir, file)
|
||
shutil.copy(src_png, dst_png)
|
||
self.logger.info(f"Archived plot: {file}")
|
||
|
||
self.tracker.mark_done(task_id_train)
|
||
|
||
else:
|
||
self.logger.info("Skipping Train (Already Done).")
|
||
# 恢复变量状态
|
||
self.current_nep_pot = os.path.join(step_dir, "nep.txt")
|
||
self.current_train_set = os.path.join(step_dir, "train.xyz")
|