nep框架重构

2025-12-09 01:15:38 +08:00
parent 19a6924a41
commit 91bdb0dab1
30 changed files with 7930 additions and 1001 deletions
--- a/src/workflow.py
+++ b/src/workflow.py
@@ -0,0 +1,131 @@
+# src/workflow.py
+import os
+import shutil
+import logging
+from src.utils import load_yaml
+from src.machine import MachineManager
+from src.steps import MDStep, SelectStep, SCFStep, TrainStep
+
+
+class Workflow:
+    def __init__(self, root_dir):
+        self.root_dir = root_dir
+
+        # 1. 加载配置
+        self.param = load_yaml(os.path.join(root_dir, "config/param.yaml"))
+
+        # 2. 初始化机器管理器
+        self.machine = MachineManager(os.path.join(root_dir, "config/machine.yaml"))
+
+        # 3. 初始化路径变量
+        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()
+
+        # 状态追踪变量
+        self.current_nep_pot = os.path.join(self.data_dir, self.param['files']['initial_pot'])
+        # 假设第一轮之前的 train set 也是空的或者由用户提供，这里先指向一个基础文件
+        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 >>> Starting Iteration: {iter_id} <<<")
+            os.makedirs(iter_path, exist_ok=True)
+
+            # --- 执行该轮定义的各个 Step ---
+            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")
+
+                    # 只有第一轮且有init需求时，才进行 POSCAR -> model.xyz 转化
+                    # 这里为了 Local 测试，我们简单处理：直接把 POSCAR 拷过去当 model.xyz (仅作演示)
+                    # 实际上你应该调用 gpumdkit 转化
+                    if iter_id == 0:
+                        os.makedirs(step_dir, exist_ok=True)
+                        shutil.copy(os.path.join(self.data_dir, self.param['files']['poscar']),
+                                    os.path.join(step_dir, "model.xyz"))
+
+                    # 遍历子任务 (preheat, production...)
+                    for sub in step_conf.get('sub_tasks', []):
+                        template_sub_name = sub['template_sub']
+                        sub_work_dir = os.path.join(step_dir, template_sub_name)
+                        template_path = os.path.join(self.template_dir, "00.md", template_sub_name)
+
+                        # 实例化并运行
+                        md_task = MDStep(f"MD-{template_sub_name}", sub_work_dir, self.machine, self.config)
+
+                        # 关键：要把上一级准备好的 model.xyz 拷进来
+                        if iter_id == 0:
+                            shutil.copy(os.path.join(step_dir, "model.xyz"), os.path.join(sub_work_dir, "model.xyz"))
+                        # 如果是后续轮次，应该用上一轮选好的结构，这里暂略，先跑通第一轮
+
+                        md_task.run(self.current_nep_pot, template_path)
+
+                        # 记录最后生成的 dump.xyz 路径，供下一步使用
+                        self.last_dump_path = os.path.join(sub_work_dir, "dump.xyz")
+
+                # ==========================
+                # Step: 01.select
+                # ==========================
+                elif step_name == "01.select":
+                    step_dir = os.path.join(iter_path, "01.select")
+                    select_task = SelectStep("Select", step_dir, self.machine, self.config)
+
+                    # 使用上一步产生的 dump 和 当前的训练集/势函数
+                    select_task.run(
+                        dump_path=getattr(self, 'last_dump_path', None),
+                        train_path=self.current_train_set,
+                        nep_path=self.current_nep_pot,
+                        method=step_conf.get('method'),
+                        params=step_conf.get('params')
+                    )
+
+                # ==========================
+                # Step: 02.scf
+                # ==========================
+                elif step_name == "02.scf":
+                    step_dir = os.path.join(iter_path, "02.scf")
+                    scf_task = SCFStep("SCF", step_dir, self.machine, self.config)
+
+                    template_path = os.path.join(self.template_dir, "02.scf")
+                    potcar_path = os.path.join(self.data_dir, self.param['files']['potcar'])
+
+                    scf_task.run(template_path, potcar_path)
+
+                    # 假装产生了一些新数据
+                    self.new_data_chunk = os.path.join(step_dir, "scf_results.xyz")
+
+                # ==========================
+                # Step: 03.train
+                # ==========================
+                elif step_name == "03.train":
+                    step_dir = os.path.join(iter_path, "03.train")
+                    train_task = TrainStep("Train", step_dir, self.machine, self.config)
+
+                    template_path = os.path.join(self.template_dir, "03.train")
+
+                    # 实际逻辑应该是把 self.new_data_chunk 合并到 total_train.xyz
+                    # 这里直接传入
+                    train_task.run(template_path, getattr(self, 'new_data_chunk', None))
+
+                    # 更新当前势函数路径，供下一轮使用
+                    self.current_nep_pot = os.path.join(step_dir, "nep.txt")
+
+        self.logger.info("Workflow Finished Successfully.")
+
+    @property
+    def config(self):
+        return self.param  # 简单透传