nep框架搭建

nep_auto/modules/m3_scf.py

@@ -0,0 +1,91 @@
+import shutil
+from pathlib import Path
+from ase.io import read, write
+from nep_auto.modules.base_module import BaseModule
+
+
+class SCFModule(BaseModule):
+    def __init__(self, driver, iter_id):
+        super().__init__(driver, iter_id)
+        self.template_subdir = "02_scf"
+        self.work_dir = self.iter_dir / "02.scf"
+        self.select_dir = self.iter_dir / "01.select"
+
+    def get_work_dir(self):
+        return self.work_dir
+
+    def run(self):
+        self.logger.info(f"⚛️ [SCF] Starting DFT Calculation Iter {self.iter_id}...")
+        self.initialize()
+
+        # 1. 读取 selected.xyz
+        selected_xyz = self.select_dir / "selected.xyz"
+        if not selected_xyz.exists():
+            raise FileNotFoundError("selected.xyz missing")
+
+        self.logger.info("   -> Reading structures using ASE...")
+        atoms_list = read(selected_xyz, index=':')
+        self.logger.info(f"   -> Found {len(atoms_list)} structures.")
+
+        # 2. 准备任务文件夹
+        task_dirs = []
+        for i, atoms in enumerate(atoms_list):
+            task_name = f"task.{i:03d}"
+            task_dir = self.work_dir / task_name
+            task_dir.mkdir(exist_ok=True)
+            task_dirs.append(task_dir)
+
+            # 写 POSCAR
+            write(task_dir / "POSCAR", atoms, format='vasp')
+
+            # 复制模版 INCAR, KPOINTS, POTCAR
+            self.copy_template("INCAR", target_name=None)  # 复制到 self.work_dir
+            shutil.copy(self.work_dir / "INCAR", task_dir / "INCAR")  # 再分发
+            self.copy_template("KPOINTS", target_name=None)
+            shutil.copy(self.work_dir / "KPOINTS", task_dir / "KPOINTS")
+            self.copy_template("POTCAR", target_name=None)
+            shutil.copy(self.work_dir / "POTCAR", task_dir / "POTCAR")
+
+        # 3. 提交任务
+        # 这里区分 local 模式和 slurm 模式
+        # 既然你目前是 interactive gpu，我们假设是串行或者简单的并行
+        self.logger.info("   -> Running VASP jobs...")
+
+        success_count = 0
+        for task_dir in task_dirs:
+            self.logger.info(f"      -> Running {task_dir.name}...")
+            try:
+                # 调用 machine.yaml 里定义的 vasp
+                # 注意：如果 task 很多，这里最好写成多进程并发
+                self.runner.run("vasp", cwd=task_dir)
+
+                # 简单检查
+                if (task_dir / "OUTCAR").exists():
+                    success_count += 1
+            except Exception as e:
+                self.logger.error(f"Task {task_dir.name} failed: {e}")
+
+        self.logger.info(f"   -> Finished. Success: {success_count}/{len(task_dirs)}")
+
+        # 4. 收集数据 (OUTCAR -> NEP-dataset.xyz)
+        self.logger.info("   -> Collecting data...")
+        valid_atoms = []
+        for task_dir in task_dirs:
+            try:
+                # 读取 OUTCAR
+                atoms = read(task_dir / "OUTCAR", format='vasp-outcar')
+                valid_atoms.append(atoms)
+            except:
+                pass
+
+        if valid_atoms:
+            write(self.work_dir / "NEP-dataset.xyz", valid_atoms, format='extxyz')
+        else:
+            raise RuntimeError("No valid OUTCARs found!")
+
+        self.check_done()
+
+    def check_done(self):
+        if (self.work_dir / "NEP-dataset.xyz").exists():
+            return True
+        raise RuntimeError("SCF failed: NEP-dataset.xyz not generated")