116 lines
5.0 KiB
Python
116 lines
5.0 KiB
Python
import shutil
|
||
import glob
|
||
from pathlib import Path
|
||
from nep_auto.modules.base_module import BaseModule
|
||
|
||
|
||
class MDModule(BaseModule):
|
||
def __init__(self, driver, iter_id):
|
||
super().__init__(driver, iter_id)
|
||
self.template_subdir = "00_md"
|
||
# 预热目录 (输入源)
|
||
self.preheat_dir = self.iter_dir / "00.md" / "preheat"
|
||
# MD 目录 (工作区)
|
||
self.work_dir = self.iter_dir / "00.md" / "md"
|
||
|
||
def get_work_dir(self):
|
||
return self.work_dir
|
||
|
||
def run(self):
|
||
self.logger.info(f"🌪️ [MD] Starting Sampling Phase Iter {self.iter_id}...")
|
||
self.initialize()
|
||
|
||
# ----------------------------------------
|
||
# 1. 从预热轨迹中采样 (dump.xyz -> sampled_structures.xyz)
|
||
# ----------------------------------------
|
||
preheat_dump = self.preheat_dir / "dump.xyz"
|
||
if not preheat_dump.exists():
|
||
raise FileNotFoundError(f"Preheat dump not found: {preheat_dump}")
|
||
|
||
# 调用 sample_structures.py
|
||
# 假设参数: input_file method number
|
||
kit_root = self.driver.config_param['env']['gpumdkit_root']
|
||
script = f"{kit_root}/Scripts/sample_structures/sample_structures.py"
|
||
|
||
# 复制 dump 到当前目录以便处理
|
||
local_dump = self.work_dir / "preheat_dump.xyz"
|
||
shutil.copy(preheat_dump, local_dump)
|
||
|
||
self.logger.info(" -> Sampling structures from preheat trajectory...")
|
||
# 按照你的描述: sample_structures.py dump.xyz uniform 4
|
||
# 这里 "4" 可以放到 param.yaml 里配置，暂时写死或读取默认
|
||
self.runner.run(
|
||
"python_script", # 这里可以用 local runner 直接跑 python
|
||
cwd=self.work_dir,
|
||
extra_args=f"{script} preheat_dump.xyz uniform 4"
|
||
)
|
||
|
||
# 产物通常叫 sampled_structures.xyz，我们需要把它作为后续 MD 的起始结构
|
||
# 但注意：GPUMD MD 通常读取 model.xyz 或者 restart。
|
||
# 如果你的 run.in 里写的是 load_xyz sampled_structures.xyz，那就没问题。
|
||
# 如果不是，通常做法是把 sampled_structures.xyz 切分成多个文件夹。
|
||
|
||
# --- 修正逻辑：根据你的描述 "生成 sample_1-4 文件夹" ---
|
||
# 我们遍历 template/00_md/md_run_*.in
|
||
tpl_path = Path("template") / self.template_subdir
|
||
run_templates = sorted(list(tpl_path.glob("md_run_*.in")))
|
||
|
||
if not run_templates:
|
||
self.logger.warning(f"⚠️ No 'md_run_*.in' found in {tpl_path}, looking for 'run.in'...")
|
||
run_templates = list(tpl_path.glob("run.in"))
|
||
|
||
sub_tasks = []
|
||
nep_source = self.preheat_dir / "nep.txt" # 沿用预热阶段的势函数
|
||
|
||
for idx, tpl in enumerate(run_templates, start=1):
|
||
task_name = f"sample_{idx}"
|
||
task_dir = self.work_dir / task_name
|
||
task_dir.mkdir(exist_ok=True)
|
||
sub_tasks.append(task_dir)
|
||
|
||
# 1. 复制 run.in
|
||
shutil.copy(tpl, task_dir / "run.in")
|
||
|
||
# 2. 复制 nep.txt
|
||
shutil.copy(nep_source, task_dir / "nep.txt")
|
||
|
||
# 3. 复制结构 (假设所有 sample 都从预热的最后一帧或 sampled_structures 开始)
|
||
# 这里简化处理：复制 model.xyz (初始结构) 或者 使用 preheat 的最后状态
|
||
# 根据你的流程，通常需要把 sampled_structures.xyz 里的某一帧放进去
|
||
# 或者 GPUMD 支持直接读取 exyz。
|
||
# 这里我们假设 run.in 里配置好了读取方式，我们只负责给文件。
|
||
if (self.preheat_dir / "model.xyz").exists():
|
||
shutil.copy(self.preheat_dir / "model.xyz", task_dir / "model.xyz")
|
||
|
||
# ----------------------------------------
|
||
# 2. 执行所有 Sample 任务
|
||
# ----------------------------------------
|
||
self.logger.info(f" -> Submitting {len(sub_tasks)} MD tasks...")
|
||
|
||
for task_dir in sub_tasks:
|
||
self.logger.info(f" -> Running {task_dir.name}...")
|
||
self.runner.run("gpumd", cwd=task_dir)
|
||
|
||
# ----------------------------------------
|
||
# 3. 合并结果
|
||
# ----------------------------------------
|
||
self.logger.info(" -> Merging dump files...")
|
||
# cat sample_*/dump.xyz >> dump.xyz
|
||
# 使用 python 实现 cat 以跨平台安全
|
||
target_dump = self.work_dir / "dump.xyz"
|
||
with open(target_dump, 'wb') as outfile:
|
||
for task_dir in sub_tasks:
|
||
src = task_dir / "dump.xyz"
|
||
if src.exists():
|
||
with open(src, 'rb') as infile:
|
||
shutil.copyfileobj(infile, outfile)
|
||
else:
|
||
self.logger.warning(f"⚠️ {task_dir.name} generated no dump.xyz")
|
||
|
||
self.check_done()
|
||
|
||
def check_done(self):
|
||
if (self.work_dir / "dump.xyz").exists():
|
||
self.logger.info("✅ MD Sampling finished.")
|
||
return True
|
||
raise RuntimeError("MD failed: dump.xyz not created.") |