{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "d7315421",
   "metadata": {},
   "source": [
    "# AtomPPGen 教程 07：完整端到端示例（以 Al 为例）\n",
    "\n",
    "本 notebook 以 **Al（Z=13）** 为例，把赝势生成流程从头到尾串联成一个“可一键 Run All”的最小闭环：\n",
    "\n",
    "1. 全电子原子解（AE）\n",
    "2. TM 伪化（构造平滑伪轨道）\n",
    "3. 半局域势反演（得到 $V_l(r)$）\n",
    "4. KB 可分离形式（得到 $V_{\\mathrm{loc}}$, $\\beta_l$, $D_l$）\n",
    "5. 可转移性验证（范数守恒 / 对数导数 / 幽灵态）\n",
    "6. 导出（JSON+NPZ）\n",
    "\n",
    "如果你已经理解 00-05 的每个环节，这个 notebook 可以作为“把知识变成产物”的工作模板。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "9ee9ed4c",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:41.242929Z",
     "iopub.status.busy": "2025-12-15T08:11:41.242840Z",
     "iopub.status.idle": "2025-12-15T08:11:41.949815Z",
     "shell.execute_reply": "2025-12-15T08:11:41.949370Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Ready\n"
     ]
    }
   ],
   "source": [
    "from pathlib import Path\n",
    "import platform\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from atomppgen import (\n",
    "    solve_ae_atom,\n",
    "    tm_pseudize,\n",
    "    invert_semilocal_potential,\n",
    "    kb_transform,\n",
    "    export_pseudopotential,\n",
    ")\n",
    "from atomppgen.validate import run_full_validation\n",
    "\n",
    "# 中文字体配置（兼容多平台）\n",
    "if platform.system() == 'Darwin':\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:\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": "8046490e",
   "metadata": {},
   "source": [
    "## Step 1：全电子原子解（AE）\n",
    "\n",
    "AE 解提供参考能级与波函数，它决定了后续伪化的目标：\n",
    "\n",
    "- 外区（$r>r_c$）必须与 AE 波函数一致\n",
    "- 内区（$r\\le r_c$）可以重写，但要满足模守恒与光滑性"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c89aaeb1",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:41.951134Z",
     "iopub.status.busy": "2025-12-15T08:11:41.951017Z",
     "iopub.status.idle": "2025-12-15T08:11:50.288896Z",
     "shell.execute_reply": "2025-12-15T08:11:50.287885Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "AE eigenvalues (Ha):\n",
      "l= 0 eps= -0.2469907526930879\n",
      "l= 1 eps= 0.0009146350242940937\n",
      "l= 2 eps= 0.00718614024367731\n"
     ]
    }
   ],
   "source": [
    "Z = 13\n",
    "xc = 'PZ81'\n",
    "lmax = 2\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",
    "print('AE eigenvalues (Ha):')\n",
    "for l in range(lmax + 1):\n",
    "    print('l=', l, 'eps=', ae.eps_by_l[l][-1])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0e26d07",
   "metadata": {},
   "source": [
    "## Step 2：TM 伪化（构造平滑伪轨道）\n",
    "\n",
    "TM 伪化的关键是“在内区抹平节点/尖峰，同时保持外区完全一致”，并用模守恒锁定散射性质。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "da18a690",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:50.291802Z",
     "iopub.status.busy": "2025-12-15T08:11:50.291657Z",
     "iopub.status.idle": "2025-12-15T08:11:50.301138Z",
     "shell.execute_reply": "2025-12-15T08:11:50.300700Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "l= 0 rc= 2.1 norm_error= 4.721090412432319e-16\n",
      "l= 1 rc= 2.2 norm_error= 5.935084483710011e-16\n",
      "l= 2 rc= 2.4 norm_error= 2.9757430941508053e-16\n"
     ]
    }
   ],
   "source": [
    "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",
    "for l, tm in tm_dict.items():\n",
    "    print('l=', l, 'rc=', tm.rc, 'norm_error=', tm.norm_error)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26149c51",
   "metadata": {},
   "source": [
    "## Step 3：势反演（得到半局域势 $V_l(r)$）\n",
    "\n",
    "直觉上：当我们“换了波函数”，就必须同步“换一个势”，使它仍是同一能量下的解；\n",
    "反演公式把这个要求变成逐点代数计算。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "121e83ff",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:50.302881Z",
     "iopub.status.busy": "2025-12-15T08:11:50.302646Z",
     "iopub.status.idle": "2025-12-15T08:11:50.420014Z",
     "shell.execute_reply": "2025-12-15T08:11:50.419031Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Inversion done: [0, 1, 2]\n"
     ]
    }
   ],
   "source": [
    "inv_dict = {l: invert_semilocal_potential(tm_dict[l], ae.r) for l in tm_dict}\n",
    "print('Inversion done:', list(inv_dict.keys()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20dc4bad",
   "metadata": {},
   "source": [
    "## Step 4：KB 可分离形式（得到 $V_{\\mathrm{loc}}$, $\\beta_l$, $D_l$）\n",
    "\n",
    "半局域势在平面波里昂贵，KB 通过秩-1 投影把非局域部分写成少量标量积，从而显著加速。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "e9f802ce",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:50.422114Z",
     "iopub.status.busy": "2025-12-15T08:11:50.421982Z",
     "iopub.status.idle": "2025-12-15T08:11:50.424689Z",
     "shell.execute_reply": "2025-12-15T08:11:50.424319Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "KB loc_channel = 2\n",
      "KB nonlocal channels = [0, 1]\n",
      "KB D_l = {0: np.float64(4.425529213473979), 1: np.float64(2.6664472703540274)}\n"
     ]
    }
   ],
   "source": [
    "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",
    "print('KB loc_channel =', kb.loc_channel)\n",
    "print('KB nonlocal channels =', sorted(kb.beta_l.keys()))\n",
    "print('KB D_l =', kb.D_l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26dc3373",
   "metadata": {},
   "source": [
    "## Step 5：可转移性验证\n",
    "\n",
    "验证不是“锦上添花”，而是赝势可用性的底线：\n",
    "\n",
    "- **范数守恒**：锁定核内电荷与散射信息\n",
    "- **对数导数**：在能量窗口内比较散射性质（相移）\n",
    "- **幽灵态**：排除非物理的深束缚态"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "22e99d36",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:50.425954Z",
     "iopub.status.busy": "2025-12-15T08:11:50.425854Z",
     "iopub.status.idle": "2025-12-15T08:11:50.595178Z",
     "shell.execute_reply": "2025-12-15T08:11:50.594714Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "overall_passed = True\n",
      "norm pass = True\n",
      "log-derivative pass = True\n",
      "ghost pass = True\n"
     ]
    }
   ],
   "source": [
    "report = run_full_validation(\n",
    "    ae,\n",
    "    tm_dict=tm_dict,\n",
    "    inv_dict=inv_dict,\n",
    "    r_test=3.0,\n",
    "    E_range_Ry=(-0.5, 0.5),\n",
    "    E_step_Ry=0.05,\n",
    ")\n",
    "\n",
    "print('overall_passed =', report.overall_passed)\n",
    "print('norm pass =', all(r.passed for r in report.norm_results.values()))\n",
    "print('log-derivative pass =', all(r.passed for r in report.log_deriv_results.values()))\n",
    "print('ghost pass =', report.ghost_result.passed)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c707e9d",
   "metadata": {},
   "source": [
    "## Step 6：导出（JSON+NPZ）\n",
    "\n",
    "导出后你将得到两类文件：\n",
    "\n",
    "- `*.json`：面向阅读/审计（参数、单位、验证报告、以及 NPZ 文件索引）\n",
    "- `*.npz`：面向数值计算（网格、半局域势、KB 投影子等数组）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "b125267a",
   "metadata": {
    "execution": {
     "iopub.execute_input": "2025-12-15T08:11:50.596968Z",
     "iopub.status.busy": "2025-12-15T08:11:50.596831Z",
     "iopub.status.idle": "2025-12-15T08:11:50.611857Z",
     "shell.execute_reply": "2025-12-15T08:11:50.610641Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[PosixPath('/Users/gilbert/Downloads/同步空间/Work/Courses/大四上/电子结构理论与计算/代码作业/2-3/AtomPPGen/outputs/tutorial_07/al_lda_complete.json'),\n",
       " PosixPath('/Users/gilbert/Downloads/同步空间/Work/Courses/大四上/电子结构理论与计算/代码作业/2-3/AtomPPGen/outputs/tutorial_07/al_lda_complete.npz')]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "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_07'\n",
    "out_dir.mkdir(parents=True, exist_ok=True)\n",
    "output_prefix = out_dir / 'al_lda_complete'\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"
   ]
  }
 ],
 "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
}