91 lines
3.4 KiB
Python
91 lines
3.4 KiB
Python
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") |