{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "80d381d6",
   "metadata": {},
   "source": [
    "# AtomPPGen 教程 06：导出赝势（JSON/NPZ）\n",
    "\n",
    "本节目标：把一套已生成并通过验证的赝势数据导出到 **JSON（元数据/验证报告）** 与 **NPZ（数值数组）** 两种格式，并解释输出字段的含义。\n",
    "\n",
    "## Colab 环境（可选）\n",
    "\n",
    "如果你在 Google Colab 运行，可先在一个 **shell 单元** 中执行：\n",
    "\n",
    "```bash\n",
    "pip install -q git+https://github.com/bud-primordium/AtomSCF.git@main\n",
    "pip install -q git+https://github.com/bud-primordium/AtomPPGen.git@main\n",
    "```\n",
    "\n",
    "本地环境建议使用可编辑安装（`pip install -e`），并确保 `AtomSCF` 已安装。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "00a3ee40",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:30.651200Z",
     "iopub.status.busy": "2025-12-15T08:11:30.651099Z",
     "iopub.status.idle": "2025-12-15T08:11:31.394007Z",
     "shell.execute_reply": "2025-12-15T08:11:31.393597Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ready\n"
     ]
    }
   ],
   "source": [
    "# 常用依赖与字体配置\n",
    "from pathlib import Path\n",
    "import json\n",
    "import platform\n",
    "\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from atomppgen import solve_ae_atom, tm_pseudize, invert_semilocal_potential, kb_transform, export_pseudopotential\n",
    "from atomppgen.validate import run_full_validation\n",
    "\n",
    "# 中文字体配置（兼容多平台）\n",
    "if platform.system() == 'Darwin':  # macOS\n",
    "    plt.rcParams['font.sans-serif'] = ['Arial Unicode MS', 'Heiti TC', 'STHeiti']\n",
    "elif platform.system() == 'Windows':\n",
    "    plt.rcParams['font.sans-serif'] = ['Microsoft YaHei', 'SimHei']\n",
    "else:  # Linux / Colab\n",
    "    plt.rcParams['font.sans-serif'] = ['DejaVu Sans', 'WenQuanYi Micro Hei']\n",
    "plt.rcParams['axes.unicode_minus'] = False\n",
    "\n",
    "print('Ready')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19ab3a76",
   "metadata": {},
   "source": [
    "## 生成一套示例赝势（Al, LDA-PZ81）\n",
    "\n",
    "导出模块需要的输入来自一条标准流水线：\n",
    "\n",
    "`solve_ae_atom → tm_pseudize → invert_semilocal_potential → kb_transform → run_full_validation → export_pseudopotential`\n",
    "\n",
    "说明：`export_pseudopotential()` 会同时写出 `*.json` 与 `*.npz`（推荐配对使用）。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "41b8b8e3",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:31.396022Z",
     "iopub.status.busy": "2025-12-15T08:11:31.395854Z",
     "iopub.status.idle": "2025-12-15T08:11:39.161421Z",
     "shell.execute_reply": "2025-12-15T08:11:39.160666Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "overall_passed = True\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "[PosixPath('/Users/gilbert/Downloads/同步空间/Work/Courses/大四上/电子结构理论与计算/代码作业/2-3/AtomPPGen/outputs/tutorial_06/al_lda.json'),\n",
       " PosixPath('/Users/gilbert/Downloads/同步空间/Work/Courses/大四上/电子结构理论与计算/代码作业/2-3/AtomPPGen/outputs/tutorial_06/al_lda.npz')]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# 1) 全电子原子解（AE）\n",
    "Z = 13\n",
    "xc = 'PZ81'\n",
    "lmax = 2  # s,p,d\n",
    "\n",
    "ae = solve_ae_atom(\n",
    "    Z=Z,\n",
    "    spin_mode='LDA',\n",
    "    xc=xc,\n",
    "    lmax=lmax,\n",
    "    grid_type='exp_transformed',\n",
    "    grid_params={'n': 800, 'rmax': 100.0},\n",
    "    scf_params={'tol': 1e-6, 'maxiter': 150},\n",
    ")\n",
    "\n",
    "# 2) TM 伪化（按通道）\n",
    "rc_by_l = {0: 2.1, 1: 2.2, 2: 2.4}\n",
    "tm_dict = {}\n",
    "for l in range(lmax + 1):\n",
    "    tm_dict[l] = tm_pseudize(\n",
    "        ae.r, ae.w,\n",
    "        ae.u_by_l[l][-1],\n",
    "        ae.eps_by_l[l][-1],\n",
    "        l=l,\n",
    "        rc=rc_by_l[l],\n",
    "    )\n",
    "\n",
    "# 3) 势反演（半局域势）\n",
    "inv_dict = {l: invert_semilocal_potential(tm_dict[l], ae.r) for l in tm_dict}\n",
    "\n",
    "# 4) KB 可分离形式（投影子与耦合系数）\n",
    "kb = kb_transform(inv_dict, {l: ae.u_by_l[l][-1] for l in tm_dict}, ae.r, ae.w, loc_channel=2)\n",
    "\n",
    "# 5) 可转移性验证\n",
    "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)\n",
    "print('overall_passed =', report.overall_passed)\n",
    "\n",
    "# 6) 导出（JSON + NPZ）\n",
    "cwd = Path.cwd().resolve()\n",
    "repo_root = next(p for p in [cwd, *cwd.parents] if (p / 'AtomPPGen' / 'pyproject.toml').is_file())\n",
    "out_dir = repo_root / 'AtomPPGen' / 'outputs' / 'tutorial_06'\n",
    "out_dir.mkdir(parents=True, exist_ok=True)\n",
    "output_prefix = out_dir / 'al_lda'\n",
    "\n",
    "files = export_pseudopotential(\n",
    "    ae_result=ae,\n",
    "    tm_dict=tm_dict,\n",
    "    inv_dict=inv_dict,\n",
    "    validation_report=report,\n",
    "    kb_result=kb,\n",
    "    output_prefix=str(output_prefix),\n",
    "    formats=['json', 'npz'],\n",
    ")\n",
    "files"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "439a3af7",
   "metadata": {},
   "source": [
    "## 输出结构解读\n",
    "\n",
    "### JSON：面向“可读性”的结构化元数据\n",
    "\n",
    "- `metadata`：元素、泛函、单位、版本信息\n",
    "- `generation_params`：生成参数（例如 `rc_by_l`）\n",
    "- `validation_report`：范数/对数导数/幽灵态与总通过状态\n",
    "- `pseudopotential.data_file`：指向对应的 `*.npz` 数值文件\n",
    "\n",
    "### NPZ：面向“数值计算”的数组容器\n",
    "\n",
    "- `radial_grid`, `radial_weights`：径向网格与积分权重\n",
    "- `semilocal_potential_l{l}`：半局域势通道数组\n",
    "- `kb_V_loc`, `kb_beta_l{l}`, `kb_D_l{l}`：KB 局域势、投影子与耦合系数\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "44417b0c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:39.164045Z",
     "iopub.status.busy": "2025-12-15T08:11:39.163909Z",
     "iopub.status.idle": "2025-12-15T08:11:39.168293Z",
     "shell.execute_reply": "2025-12-15T08:11:39.167845Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "JSON top-level keys: ['metadata', 'generation_params', 'all_electron_reference', 'pseudopotential', 'validation_report', 'radial_grid']\n",
      "JSON metadata: dict_keys(['element', 'xc', 'generation_date', 'code', 'units'])\n",
      "JSON pp keys: dict_keys(['channels', 'data_file', 'note', 'kb'])\n",
      "NPZ key count: 22\n",
      "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']\n"
     ]
    }
   ],
   "source": [
    "# 读取并快速浏览 JSON/NPZ\n",
    "json_path = Path(str(output_prefix) + '.json')\n",
    "npz_path = Path(str(output_prefix) + '.npz')\n",
    "\n",
    "with json_path.open('r', encoding='utf-8') as f:\n",
    "    data = json.load(f)\n",
    "\n",
    "print('JSON top-level keys:', list(data.keys()))\n",
    "print('JSON metadata:', data['metadata'].keys())\n",
    "print('JSON pp keys:', data['pseudopotential'].keys())\n",
    "\n",
    "npz = np.load(npz_path)\n",
    "keys = sorted(npz.files)\n",
    "print('NPZ key count:', len(keys))\n",
    "print('NPZ sample keys:', keys[:20])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "202f92cb",
   "metadata": {},
   "source": [
    "## 关于 UPF（Quantum ESPRESSO）\n",
    "\n",
    "当前 AtomPPGen 的导出接口以 JSON/NPZ 为主，UPF 导出属于后续扩展。\n",
    "如果你需要直接用于 QE，请优先确认：\n",
    "\n",
    "1. 你的目标程序需要的是 **半局域形式** 还是 **KB 可分离形式**；\n",
    "2. 需要哪些信息（局域势、非局域投影子、参考能量、径向网格等）。\n",
    "\n",
    "这些信息已经在 JSON/NPZ 中明确保存，便于后续补齐 UPF writer。"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}