AtomPPGen 教程 06：导出赝势（JSON/NPZ）

本节目标：把一套已生成并通过验证的赝势数据导出到 JSON（元数据/验证报告） 与 NPZ（数值数组） 两种格式，并解释输出字段的含义。

Colab 环境（可选）

如果你在 Google Colab 运行，可先在一个 shell 单元 中执行：

pip install -q git+https://github.com/bud-primordium/AtomSCF.git@main
pip install -q git+https://github.com/bud-primordium/AtomPPGen.git@main

本地环境建议使用可编辑安装（pip install -e），并确保 AtomSCF 已安装。

[1]:

# 常用依赖与字体配置
from pathlib import Path
import json
import platform

import numpy as np
import matplotlib.pyplot as plt

from atomppgen import solve_ae_atom, tm_pseudize, invert_semilocal_potential, kb_transform, export_pseudopotential
from atomppgen.validate import run_full_validation

# 中文字体配置（兼容多平台）
if platform.system() == 'Darwin':  # macOS
    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'Heiti TC', 'STHeiti']
elif platform.system() == 'Windows':
    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei']
else:  # Linux / Colab
    plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'WenQuanYi Micro Hei']
plt.rcParams['axes.unicode_minus'] = False

print('Ready')

Ready

生成一套示例赝势（Al, LDA-PZ81）

导出模块需要的输入来自一条标准流水线：

solve_ae_atom → tm_pseudize → invert_semilocal_potential → kb_transform → run_full_validation → export_pseudopotential

说明：export_pseudopotential() 会同时写出 *.json 与 *.npz（推荐配对使用）。

[2]:

# 1) 全电子原子解（AE）
Z = 13
xc = 'PZ81'
lmax = 2  # s,p,d

ae = solve_ae_atom(
    Z=Z,
    spin_mode='LDA',
    xc=xc,
    lmax=lmax,
    grid_type='exp_transformed',
    grid_params={'n': 800, 'rmax': 100.0},
    scf_params={'tol': 1e-6, 'maxiter': 150},
)

# 2) TM 伪化（按通道）
rc_by_l = {0: 2.1, 1: 2.2, 2: 2.4}
tm_dict = {}
for l in range(lmax + 1):
    tm_dict[l] = tm_pseudize(
        ae.r, ae.w,
        ae.u_by_l[l][-1],
        ae.eps_by_l[l][-1],
        l=l,
        rc=rc_by_l[l],
    )

# 3) 势反演（半局域势）
inv_dict = {l: invert_semilocal_potential(tm_dict[l], ae.r) for l in tm_dict}

# 4) KB 可分离形式（投影子与耦合系数）
kb = kb_transform(inv_dict, {l: ae.u_by_l[l][-1] for l in tm_dict}, ae.r, ae.w, loc_channel=2)

# 5) 可转移性验证
report = run_full_validation(ae, tm_dict=tm_dict, inv_dict=inv_dict, r_test=3.0, E_range_Ry=(-0.5, 0.5), E_step_Ry=0.05)
print('overall_passed =', report.overall_passed)

# 6) 导出（JSON + NPZ）
cwd = Path.cwd().resolve()
repo_root = next(p for p in [cwd, *cwd.parents] if (p / 'AtomPPGen' / 'pyproject.toml').is_file())
out_dir = repo_root / 'AtomPPGen' / 'outputs' / 'tutorial_06'
out_dir.mkdir(parents=True, exist_ok=True)
output_prefix = out_dir / 'al_lda'

files = export_pseudopotential(
    ae_result=ae,
    tm_dict=tm_dict,
    inv_dict=inv_dict,
    validation_report=report,
    kb_result=kb,
    output_prefix=str(output_prefix),
    formats=['json', 'npz'],
)
files

overall_passed = True

[2]:

[PosixPath('/Users/gilbert/Downloads/同步空间/Work/Courses/大四上/电子结构理论与计算/代码作业/2-3/AtomPPGen/outputs/tutorial_06/al_lda.json'),
 PosixPath('/Users/gilbert/Downloads/同步空间/Work/Courses/大四上/电子结构理论与计算/代码作业/2-3/AtomPPGen/outputs/tutorial_06/al_lda.npz')]

输出结构解读

JSON：面向“可读性”的结构化元数据

metadata：元素、泛函、单位、版本信息
generation_params：生成参数（例如 rc_by_l）
validation_report：范数/对数导数/幽灵态与总通过状态
pseudopotential.data_file：指向对应的 *.npz 数值文件

NPZ：面向“数值计算”的数组容器

radial_grid, radial_weights：径向网格与积分权重
semilocal_potential_l{l}：半局域势通道数组
kb_V_loc, kb_beta_l{l}, kb_D_l{l}：KB 局域势、投影子与耦合系数

[3]:

# 读取并快速浏览 JSON/NPZ
json_path = Path(str(output_prefix) + '.json')
npz_path = Path(str(output_prefix) + '.npz')

with json_path.open('r', encoding='utf-8') as f:
    data = json.load(f)

print('JSON top-level keys:', list(data.keys()))
print('JSON metadata:', data['metadata'].keys())
print('JSON pp keys:', data['pseudopotential'].keys())

npz = np.load(npz_path)
keys = sorted(npz.files)
print('NPZ key count:', len(keys))
print('NPZ sample keys:', keys[:20])

JSON top-level keys: ['metadata', 'generation_params', 'all_electron_reference', 'pseudopotential', 'validation_report', 'radial_grid']
JSON metadata: dict_keys(['element', 'xc', 'generation_date', 'code', 'units'])
JSON pp keys: dict_keys(['channels', 'data_file', 'note', 'kb'])
NPZ key count: 22
NPZ sample keys: ['Z', 'ae_eigenvalues_l0', 'ae_eigenvalues_l1', 'ae_eigenvalues_l2', 'ae_wavefunction_l0_n3', 'ae_wavefunction_l1_n3', 'ae_wavefunction_l2_n3', 'kb_D_l0', 'kb_D_l1', 'kb_V_loc', 'kb_beta_l0', 'kb_beta_l1', 'kb_loc_channel', 'ps_wavefunction_l0', 'ps_wavefunction_l1', 'ps_wavefunction_l2', 'radial_grid', 'radial_weights', 'semilocal_potential_l0', 'semilocal_potential_l1']

关于 UPF（Quantum ESPRESSO）

当前 AtomPPGen 的导出接口以 JSON/NPZ 为主，UPF 导出属于后续扩展。如果你需要直接用于 QE，请优先确认：

你的目标程序需要的是 半局域形式 还是 KB 可分离形式；
需要哪些信息（局域势、非局域投影子、参考能量、径向网格等）。

这些信息已经在 JSON/NPZ 中明确保存，便于后续补齐 UPF writer。