使用 abipy 批量提取 VASP 几何优化的最终结构

March 19, 2026
Published in 计算材料学

Abstract

在进行密度泛函理论(DFT)等高通量计算任务时,从大量的计算目录下提取出几何优化完成的最终晶体结构往往是一项繁琐的工作。本文将分享一个实用的 Python 脚本,它整合了 abipypymatgen,能够自动遍历计算目录、判断收敛状态,并高效地提取出用于后续计算的 VASP POSCAR 文件。

Keywords: Python, VASP, abipy, 数据处理, 自动化脚本, 计算晶体学

脚本功能概述

在材料模拟与计算化学领域,结构优化是获取系统基态结构的关键步骤。当涉及多组超胞成分或者大量不同掺杂构型的计算时,手动一个个打开文件提取最终结构的效率极低。本脚本的主要功能包括:

  1. 自动查找计算目录:根据指定的路径和通配符模式(Pattern),支持以末尾数字顺序自动识别各计算文件夹。
  2. 状态检查与收敛过滤:自动读取 run.abo 输出文件的尾部信息来判断任务状态(如 DONE、RUNNING 或是 ERROR),默认只处理已经成功收敛完成(DONE)的计算。
  3. 结构提取与格式转换:调用 abipy 库读取 HIST.nc 文件,获取最后一个离子步的结构对象,并借助 pymatgen 转化为标准 VASP POSCAR (.vasp) 格式。
  4. 批量处理与容错机制:支持覆盖现有文件(--force)与模拟执行(--dry-run),能够较好地应对中断和重复提取问题。

工作流程解析

整个脚本的核心逻辑遵循标准的提取与写出流程,确保数据的正确性和不重复处理。以下是该脚本执行的流程图:

流程图

正如流程图所示,从开始解析命令行参数后,系统会创建输出目录并定位所有的目标匹配目录,随后逐一遍历。在每一轮循环中,如果发现尚未收敛或者缺少历史输出文件,脚本会贴心地给出跳过(SKIP)提示,保证程序高效运行。

核心代码实现

本脚本兼容了 scipy 最新版本的部分废弃 API,并提供丰富的命令行调整选项。以下为完整的 Python 源码:

#!/usr/bin/env python3
"""
利用 abipy 从已完成的几何优化计算中提取最终优化结构,
并导出为 VASP POSCAR 格式 (.vasp),保存到 structure/opt/ 目录。

输出文件命名规则:
    计算目录名 + 'o' + '.vasp'
    例:al27si1o1s-1  →  structure/opt/al27si1o1s-1o.vasp

Usage (在 conda dft 环境中运行):
    python extract_opt_structures.py [options]

Options:
    --calc-dir   根计算目录 (default: ../calculation)
    --out-dir    输出目录 (default: ../structure/opt)
    --pattern    子目录 glob pattern (default: al27si1o1s-*)
    --force      强制覆盖已存在的 .vasp 文件
    --dry-run    只打印将要处理的文件,不实际写出
    --done-only  只处理状态为 DONE 的计算(默认行为)
    --all        处理所有含 HIST.nc 的计算(包括未收敛的)
"""

import os
import re
import glob
import argparse

# ---------- scipy 兼容补丁(cumtrapz/simps 在 scipy>=1.14 中已移除)----------
import scipy.integrate as _sci_int
if not hasattr(_sci_int, 'cumtrapz'):
    _sci_int.cumtrapz = _sci_int.cumulative_trapezoid
if not hasattr(_sci_int, 'simps'):
    _sci_int.simps = _sci_int.simpson

# ---------- abipy 导入 ----------
try:
    import abipy.abilab as abilab
    ABIPY_OK = True
except ImportError:
    ABIPY_OK = False

if not ABIPY_OK:
    raise ImportError(
        "abipy 未安装或未激活 conda dft 环境。\n"
        "请先执行: conda activate dft"
    )


def find_calc_dirs(calc_root: str, pattern: str) -> list:
    """按末尾数字顺序返回匹配的计算目录。"""
    dirs = glob.glob(os.path.join(calc_root, pattern))

    def sort_key(d):
        m = re.search(r'-(\d+)$', os.path.basename(d))
        return int(m.group(1)) if m else 0

    return sorted(dirs, key=sort_key)


def check_convergence_from_abo(abo_path: str) -> str:
    """从 .abo 文件尾部判断计算状态。"""
    if not os.path.isfile(abo_path):
        return 'NO_ABO'
    try:
        with open(abo_path, 'r', errors='ignore') as f:
            tail = f.read()[-8000:]
        if 'Voluntary context switches' in tail or 'Calculation completed' in tail:
            return 'DONE'
        if 'error' in tail.lower() or 'ERROR' in tail:
            return 'ERROR'
        return 'RUNNING'
    except Exception:
        return 'UNKNOWN'


def extract_final_structure(hist_path: str, struct_name: str):
    """
    用 abipy 打开 HIST.nc,返回最终步骤的 pymatgen Structure。
    struct_name 用于设置 Structure 的 comment/label。
    """
    hist = abilab.abiopen(hist_path)
    # final_structure 是最后一个离子步的 pymatgen Structure
    structure = hist.final_structure
    hist.close()

    # 更新结构的 comment(会写入 POSCAR 第二行)
    try:
        structure.comment = struct_name
    except AttributeError:
        pass  # 某些版本的 pymatgen 不支持直接设置 comment

    return structure


def write_poscar(structure, output_path: str, struct_name: str):
    """将 pymatgen Structure 写出为 VASP POSCAR 格式。"""
    try:
        from pymatgen.io.vasp import Poscar
        poscar = Poscar(structure, comment=struct_name)
        poscar_str = poscar.get_str()
        with open(output_path, 'w') as f:
            f.write(poscar_str)
    except ImportError:
        # 降级:使用 structure.to()
        structure.to(fmt='poscar', filename=output_path)
        # 替换 POSCAR 第一行为 struct_name
        with open(output_path, 'r') as f:
            lines = f.readlines()
        if lines:
            lines[0] = struct_name + '\n'
        with open(output_path, 'w') as f:
            f.writelines(lines)


def main():
    parser = argparse.ArgumentParser(
        description='利用 abipy 提取几何优化最终结构并导出为 .vasp 文件')
    parser.add_argument('--calc-dir', default='../calculation',
                        help='根计算目录 (default: ../calculation)')
    parser.add_argument('--out-dir', default='../structure/opt',
                        help='输出目录 (default: ../structure/opt)')
    parser.add_argument('--pattern', default='al27si1o1s-*',
                        help='子目录 glob pattern (default: al27si1o1s-*)')
    parser.add_argument('--force', action='store_true',
                        help='强制覆盖已存在的 .vasp 文件')
    parser.add_argument('--dry-run', action='store_true',
                        help='只打印待处理信息,不写出文件')
    parser.add_argument('--all', dest='process_all', action='store_true',
                        help='处理所有含 HIST.nc 的目录(不限于 DONE 状态)')
    args = parser.parse_args()

    script_dir = os.path.dirname(os.path.abspath(__file__))

    # 解析路径
    if os.path.isabs(args.calc_dir):
        calc_root = args.calc_dir
    else:
        calc_root = os.path.normpath(os.path.join(script_dir, args.calc_dir))

    if os.path.isabs(args.out_dir):
        out_dir = args.out_dir
    else:
        out_dir = os.path.normpath(os.path.join(script_dir, args.out_dir))

    if not args.dry_run:
        os.makedirs(out_dir, exist_ok=True)

    print(f"计算根目录  : {calc_root}")
    print(f"输出目录    : {out_dir}")
    print(f"Pattern     : {args.pattern}")
    print(f"覆盖模式    : {'是' if args.force else '否'}")
    print(f"Dry-run     : {'是' if args.dry_run else '否'}")
    print()

    calc_dirs = find_calc_dirs(calc_root, args.pattern)
    if not calc_dirs:
        print(f"[ERROR] 未找到匹配 '{args.pattern}' 的目录:{calc_root}")
        return

    n_ok = 0
    n_skip = 0
    n_fail = 0

    for d in calc_dirs:
        dir_name = os.path.basename(d)                # e.g. al27si1o1s-1
        struct_name = dir_name + 'o'                   # e.g. al27si1o1s-1o
        out_filename = struct_name + '.vasp'           # e.g. al27si1o1s-1o.vasp
        output_path = os.path.join(out_dir, out_filename)

        hist_path = os.path.join(d, 'runo_HIST.nc')
        abo_path = os.path.join(d, 'run.abo')
        status = check_convergence_from_abo(abo_path)

        # 检查 HIST.nc 是否存在
        if not os.path.isfile(hist_path):
            print(f"  [SKIP] {dir_name:<25}  status={status}  (无 HIST.nc)")
            n_skip += 1
            continue

        # 未完成的计算默认跳过(除非 --all)
        if not args.process_all and status not in ('DONE',):
            print(f"  [SKIP] {dir_name:<25}  status={status}  (未完成,跳过)")
            n_skip += 1
            continue

        # 文件已存在且不强制覆盖
        if os.path.isfile(output_path) and not args.force:
            print(f"  [SKIP] {dir_name:<25}  status={status}  (已存在: {out_filename})")
            n_skip += 1
            continue

        if args.dry_run:
            print(f"  [DRY ] {dir_name:<25}  status={status}{out_filename}")
            n_ok += 1
            continue

        # 提取并写出
        try:
            structure = extract_final_structure(hist_path, struct_name)
            # 打印提取到的结构基本信息
            natom = len(structure)
            formula = structure.composition.reduced_formula
            vol = structure.volume
            write_poscar(structure, output_path, struct_name)
            print(f"  [OK  ] {dir_name:<25}  status={status}  "
                  f"natom={natom}  formula={formula}  vol={vol:.2f} ų  → {out_filename}")
            n_ok += 1
        except Exception as e:
            print(f"  [FAIL] {dir_name:<25}  status={status}  错误: {e}")
            n_fail += 1

    print()
    print(f"完成: 导出 {n_ok} 个,跳过 {n_skip} 个,失败 {n_fail} 个。")
    if not args.dry_run and n_ok > 0:
        print(f"文件已保存在: {out_dir}")


if __name__ == '__main__':
    main()

命令与参数说明

在此代码中,我们使用了 argparse 构建命令行接口。主要参数用法如下:

  1. --calc-dir: 指定根计算目录(默认路径 ../calculation)。
  2. --out-dir: 提取出的极小化收敛结构存放的目标位置(默认路径 ../structure/opt)。
  3. --pattern: 计算子目录的匹配表达式(例如 al27si1o1s-*)。
  4. --force: 开启此选项时,即使指定的 .vasp 文件已经存在也会被安全覆盖。
  5. --dry-run: 打印各个目录的状态以及打算执行的操作,但不真正写出文件。
  6. --all: 处理所有包含 HIST.nc 的目录,无论其收敛状态如何。

总结

利用上述自动化程序代码可以极大提升理论计算后期数据处理的效率,将从事材料计算的研究人员从繁重的文件整理中解放出来,使其能更加专注于物理机理与化学性质的分析。

Reference

  1. Ong, Shyue Ping, et al. "Python Materials Genomics (pymatgen): A robust, open-source python library for materials analysis." Computational Materials Science, vol. 68, 2013, pp. 314-319. link
  2. "abipy: Abinit python tools." Abinit Project. link