{ "cells": [ { "cell_type": "markdown", "id": "43373369", "metadata": {}, "source": [ "# Wavefunction\n", "\n", "This example shows a minimal Python workflow for wavefunctions in an atomic-orbital(AO) basis. \n", "\n", "Recall that the real-space Bloch wavefunction and the AO wavefunction are related by\n", "\n", "$$\n", "\\psi_{n \\boldsymbol{k}}(\\boldsymbol{r}) = \\frac{1}{\\sqrt{N}}\\sum_{\\boldsymbol{R}}\\sum_{j\\beta \\in \\text{uc}} e^{i\\boldsymbol{k}\\cdot \\boldsymbol{R}} c_{n, j\\beta}(\\boldsymbol{k}) \\phi_{j\\beta}(\\boldsymbol{r}-\\boldsymbol{R}-\\mathcal{R}_j),\n", "$$\n", "\n", "and the periodic part of the Bloch wavefunction and the AO wavefunction are related by\n", "\n", "$$\n", "u_{n \\boldsymbol{k}}(\\boldsymbol{r}) = \\frac{1}{\\sqrt{N}}\\sum_{\\boldsymbol{R}}\\sum_{j\\beta \\in \\text{uc}} e^{-i\\boldsymbol{k}\\cdot (\\boldsymbol{r}-\\boldsymbol{R})} c_{n, j\\beta}(\\boldsymbol{k}) \\phi_{j\\beta}(\\boldsymbol{r}-\\boldsymbol{R}-\\mathcal{R}_j),\n", "$$\n", "\n", "here $\\phi_{j\\beta}$ are the basis functions, and $c_{n, j\\beta}(\\boldsymbol{k})$ are the AO wavefunction coefficients.\n", "\n", "Starting from a DeepH data point containing `hamiltonian.h5`, `overlap.h5`, `POSCAR`, and `info.json`, we will:\n", "\n", "1. diagonalize the generalized eigenvalue problem $H(k)c = E S(k)c$ at the Gamma point,\n", "2. save the AO wavefunction coefficients to `wavefunction_ao.h5`, and\n", "3. use `AOWfnObj` to evaluate and plot the Gamma-point highest valence-band wavefunction in real space. (currently only supports SIESTA and OpenMX.)" ] }, { "cell_type": "markdown", "id": "6af8e328", "metadata": {}, "source": [ "## Setup" ] }, { "cell_type": "code", "execution_count": 1, "id": "8b0c2c30", "metadata": {}, "outputs": [], "source": [ "import os\n", "from pathlib import Path\n", "\n", "# Use writable cache directories on shared systems.\n", "os.environ.setdefault(\"MPLCONFIGDIR\", \"/tmp/deepx_dock_matplotlib\")\n", "os.environ.setdefault(\"XDG_CACHE_HOME\", \"/tmp/deepx_dock_cache\")\n", "os.environ.setdefault(\"MPICH_GPU_SUPPORT_ENABLED\", \"0\")\n", "import h5py\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "from deepx_dock.CONSTANT import DEEPX_WFNAO_FILENAME\n", "from deepx_dock.compute.eigen.hamiltonian import HamiltonianObj\n", "from deepx_dock.compute.eigen.wavefunction import AOWfnObj\n", "\n", "DATA_ROOT = Path(\"wavefunction.clean/BLBP\")\n", "DATA_DIR = DATA_ROOT / \"0\"\n", "BASIS_DIR = DATA_ROOT / \"basis\"\n", "FIG_DIR = Path(\"figures.bak\")\n", "FIG_DIR.mkdir(exist_ok=True)\n", "\n", "E_FERMI_EV = -2.4558\n", "KPTS = np.array([[0.0, 0.0, 0.0]])\n", "# Real-space grid used only for this small plotting example.\n", "GRID = np.array([48, 68, 100])" ] }, { "cell_type": "markdown", "id": "07f39762", "metadata": {}, "source": [ "## 1. Diagonalize the Hamiltonian\n", "\n", "`HamiltonianObj` loads both `hamiltonian.h5` and `overlap.h5`. Calling `diag(..., bands_only=False)` solves the non-orthogonal generalized eigenvalue problem and returns both energies and eigenvectors.\n", "\n", "- `eigvals` has shape `(n_band, n_k)`.\n", "- `eigvecs` has shape `(n_orb, n_band, n_k)`." ] }, { "cell_type": "code", "execution_count": 2, "id": "f5ca2f41", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Number of orbitals: 1584\n", "Eigenvalues shape: (1584, 1)\n", "Eigenvectors shape: (1584, 1584, 1)\n", "Gamma VBM band index: 179\n", "Gamma VBM energy: -2.626961019206 eV\n", "Gamma CBM band index: 180\n", "Gamma CBM energy: -2.186534744865 eV\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r\n", "Diagonalizing: 0%| | 0/1 [00:00 (n_k, n_band, n_orb).\n", "wfnao = eigvecs.transpose(2, 1, 0)\n", "\n", "wfn.load(\n", " KPTS,\n", " wfnao,\n", " eigvals.T,\n", " spinful=False,\n", " efermi=E_FERMI_EV,\n", " kgrid=(1, 1, 1),\n", ")\n", "\n", "wfn.to_h5()\n", "wfnao_path = DATA_DIR / DEEPX_WFNAO_FILENAME\n", "\n", "print(f\"Saved AO wavefunction coefficients to: {wfnao_path}\")\n", "with h5py.File(wfnao_path, \"r\") as h5f:\n", " print(\"Saved datasets:\", sorted(h5f.keys()))\n", " print(\"eigenvectors shape:\", h5f[\"eigenvectors\"].shape)" ] }, { "cell_type": "markdown", "id": "e5575d1e", "metadata": {}, "source": [ "## 3. Compute and plot the real-space wavefunction\n", "\n", "`AOWfnObj.to_real_space` evaluates the selected AO wavefunction on a uniform real-space grid. Here we evaluate only the Gamma-point highest valence band.\n", "\n", "If the wavefunction is evaluated on a real space grid of size $n_x \\times n_y \\times n_z$, the real-space wavefunction is normalized to $\\sqrt{n_x n_y n_z}$." ] }, { "cell_type": "code", "execution_count": 4, "id": "78879b39", "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Real-space wavefunction shape: (1, 1, 1, 48, 68, 100)\n", "Norm: 1.000010331947\n" ] } ], "source": [ "# psi shape: (selected_n_k, selected_n_band, n_spinor, nx, ny, nz).\n", "psi = wfn.to_real_space(\n", " ik=0,\n", " ib=vbm,\n", " gridsize=GRID,\n", " return_periodic=True,\n", ")\n", "\n", "psi_vbm = psi[0, 0, 0]\n", "norm = np.sum(np.abs(psi_vbm) ** 2)\n", "\n", "print(f\"Real-space wavefunction shape: {psi.shape}\")\n", "print(f\"Norm: {norm/np.prod(GRID):.12f}\")" ] }, { "cell_type": "code", "execution_count": null, "id": "d0803899", "metadata": {}, "outputs": [], "source": [ "# Use the z slice with the largest integrated density for a readable 2D image.\n", "density = np.abs(psi_vbm) ** 2\n", "z_slice = int(np.argmax(density.sum(axis=(0, 1))))\n", "fig, axes = plt.subplots(1, 2, figsize=(9, 3.8), constrained_layout=True)\n", "\n", "real_slice = psi_vbm[:, :, z_slice].real.T\n", "density_slice = density[:, :, z_slice].T\n", "\n", "# Convert grid indices to real lengths along lattice vectors a, b, and c.\n", "cell_lengths = np.linalg.norm(ham.lattice, axis=1)\n", "extent_xy = [0.0, cell_lengths[0], 0.0, cell_lengths[1]]\n", "z_coord = z_slice / GRID[2] * cell_lengths[2]\n", "\n", "real_limit = float(np.max(np.abs(real_slice)))\n", "im = axes[0].imshow(\n", " real_slice,\n", " origin=\"lower\",\n", " cmap=\"RdBu_r\",\n", " vmin=-real_limit,\n", " vmax=real_limit,\n", " aspect=\"auto\",\n", " extent=extent_xy,\n", ")\n", "axes[0].set_title(r\"$\\mathrm{Re}\\,\\psi$\")\n", "axes[0].set_xlabel(\"x (Angstrom)\")\n", "axes[0].set_ylabel(\"y (Angstrom)\")\n", "fig.colorbar(im, ax=axes[0], fraction=0.046, pad=0.04)\n", "\n", "im = axes[1].imshow(\n", " density_slice,\n", " origin=\"lower\",\n", " cmap=\"magma\",\n", " aspect=\"auto\",\n", " extent=extent_xy,\n", ")\n", "axes[1].set_title(r\"$|\\psi|^2$\")\n", "axes[1].set_xlabel(\"x (Angstrom)\")\n", "axes[1].set_ylabel(\"y (Angstrom)\")\n", "fig.colorbar(im, ax=axes[1], fraction=0.046, pad=0.04)\n", "\n", "plot_path = FIG_DIR / \"vbm.png\"\n", "fig.suptitle(f\"Gamma VBM band {vbm}, E={gamma_eigs[vbm]:.6f} eV, z={z_coord:.3f} Angstrom\")\n", "fig.savefig(plot_path, dpi=220)\n", "\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "c25d467c", "metadata": {}, "source": [ "![Gamma VBM wavefunction](figures.bak/vbm.png)" ] } ], "metadata": { "kernelspec": { "display_name": "base", "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.11" } }, "nbformat": 4, "nbformat_minor": 5 }