{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir(\"siesta_1\")\n",
    "import numpy as np\n",
    "import sisl as si\n",
    "from sisl.viz import merge_plots\n",
    "from sisl.viz.processors.math import normalize\n",
    "from functools import partial\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Siesta --- the H2O molecule\n",
    "\n",
    "This tutorial will describe a complete walk-through of some of the `sisl` functionalities that may be related to the [Siesta code](https://gitlab.com/siesta-project/siesta).\n",
    "\n",
    "## Creating the geometry\n",
    "\n",
    "Our system of interest will be the $\\mathrm H_2\\mathrm O$ system. The first task will be to create the molecule geometry.\n",
    "This is done using lists of atomic coordinates and atomic species. Additionally one needs to define the lattice (or if you prefer: unit-cell) where the molecule resides in. Siesta is a periodic DFT code and thus all directions are periodic. I.e. when simulating molecules it is vital to have a large vacuum gap between periodic images. In this case we use a supercell of side-lengths $10\\mathrm{Ang}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2o = si.Geometry(\n",
    "    [[0, 0, 0], [0.8, 0.6, 0], [-0.8, 0.6, 0.0]],\n",
    "    [si.Atom(\"O\"), si.Atom(\"H\"), si.Atom(\"H\")],\n",
    "    lattice=si.Lattice(10, origin=[-5] * 3),\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The input are the 1) xyz coordinates, 2) the atomic species and 3) the lattice that is attached.\n",
    "\n",
    "By printing the object one gets basic information regarding the geometry, such as 1) number of atoms, 2) species of atoms, 3) number of orbitals, 4) orbitals associated with each atom and 5) number of supercells (should be 1 for molecules)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(h2o)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So there are 3 atoms, 1 Oxygen and 2 Hydrogen. Currently there are only 1 orbital per atom.  \n",
    "Later we will look into the details of *orbitals* associated with *atoms* and how they may be used for wavefunctions etc.\n",
    "\n",
    "Lets visualize the atomic positions (here adding atomic indices)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "h2o.plot(axes=\"xy\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we need to create the input fdf file for Siesta:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "open(\"RUN.fdf\", \"w\").write(\n",
    "    \"\"\"%include STRUCT.fdf\n",
    "SystemLabel siesta_1\n",
    "PAO.BasisSize SZP\n",
    "MeshCutoff 250. Ry\n",
    "CDF.Save true\n",
    "CDF.Compress 9\n",
    "SaveHS true\n",
    "SaveRho true\n",
    "\"\"\"\n",
    ")\n",
    "h2o.write(\"STRUCT.fdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The first block of code simply writes a text-file with the required input, the last line of code tells the geometry (`h2o`) to write its information to the file `STRUCT.fdf`.  \n",
    "It automatically writes the following geometry information to the fdf file:\n",
    "- `LatticeConstant`\n",
    "- `LatticeVectors`\n",
    "- `NumberOfAtoms`\n",
    "- `AtomicCoordinatesFormat`\n",
    "- `AtomicCoordinatesAndAtomicSpecies`\n",
    "- `NumberOfSpecies`\n",
    "- `ChemicalSpeciesLabel`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the electronic structure\n",
    "\n",
    "Before continuing we need to run Siesta to calculate the electronic structure.\n",
    "```bash\n",
    "siesta RUN.fdf\n",
    "```\n",
    "After having completed the Siesta run we may read the Siesta output to manipulate and extract different information."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf = si.get_sile(\"RUN.fdf\")\n",
    "H = fdf.read_hamiltonian()\n",
    "# Create a short-hand to handle the geometry\n",
    "h2o = H.geometry\n",
    "print(H)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A lot of new information has appeared. The Hamiltonian object describes the non-orthogonal basis and the \"hopping\" elements between the orbitals. We see it is a non-orthogonal basis via: `orthogonal: False`. Secondly, we see it was an un-polarized calculation (`Spin{unpolarized...`).  \n",
    "Lastly the geometry information is printed again. Contrary to the previous printing of the geometry we now find additional information based on the orbitals on each of the atoms. This information is read from the Siesta output and thus the basic information regarding the orbital symmetry and the basis functions are now handled by `sisl`. The oxygen has 9 orbitals ($s+p+d$ where the $d$ orbitals are $p$-polarizations denoted by capital `P`). We also see that it is a single-$\\zeta$ calculation `Z1` (for double-$\\zeta$ `Z2` would also appear in the list). The hydrogens only has 4 orbitals $s+p$.\n",
    "For each orbital one can see its maximal radial part and how initial charges are distributed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Plotting orbitals\n",
    "\n",
    "Often it may be educational (and fun) to plot the orbital wavefunctions. To do this we use the intrinsic method in the `Orbital` class named `toGrid`. The below code is rather complicated, but the complexity is simply because we want to show the orbitals in a rectangular grid of plots."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def plot_atom(atom):\n",
    "    no = len(atom)  # number of orbitals\n",
    "    nx = no // 4\n",
    "    ny = no // nx\n",
    "    if nx * ny < no:\n",
    "        nx += 1\n",
    "    fig, axs = plt.subplots(nx, ny, figsize=(20, 5 * nx))\n",
    "    fig.suptitle(\"Atom: {}\".format(atom.symbol), fontsize=14)\n",
    "\n",
    "    def my_plot(i, orb):\n",
    "        grid = orb.toGrid(atom=atom)\n",
    "        # In case you want to plot it with external software\n",
    "        # Just uncomment the following line:\n",
    "        # grid.write(\"{}_{}.cube\".format(atom.symbol, orb.name()))\n",
    "        c, r = i // 4, (i - 4) % 4\n",
    "        if nx == 1:\n",
    "            ax = axs[r]\n",
    "        else:\n",
    "            ax = axs[c][r]\n",
    "        ax.imshow(grid.grid[:, :, grid.shape[2] // 2])\n",
    "        ax.set_title(r\"${}$\".format(orb.name(True)))\n",
    "        ax.set_xlabel(r\"$x$ [Ang]\")\n",
    "        ax.set_ylabel(r\"$y$ [Ang]\")\n",
    "\n",
    "    i = 0\n",
    "    for orb in atom:\n",
    "        my_plot(i, orb)\n",
    "        i += 1\n",
    "    if i < nx * ny:\n",
    "        # This removes the empty plots\n",
    "        for j in range(i, nx * ny):\n",
    "            c, r = j // 4, (j - 4) % 4\n",
    "            if nx == 1:\n",
    "                ax = axs[r]\n",
    "            else:\n",
    "                ax = axs[c][r]\n",
    "            fig.delaxes(ax)\n",
    "        plt.draw()\n",
    "\n",
    "\n",
    "plot_atom(h2o.atoms[0])\n",
    "plot_atom(h2o.atoms[1])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hamiltonian eigenstates\n",
    "\n",
    "At this point we have the full Hamiltonian as well as the basis functions used in the Siesta calculation.\n",
    "This completes what is needed to calculate a great deal of physical quantities, e.g. eigenstates, density of states, projected density of states and wavefunctions.\n",
    "\n",
    "To begin with we calculate the $\\Gamma$-point eigenstates and plot a subset of the eigenstates' norm on the geometry. An important aspect of the electronic structure handled by Siesta is that it is shifted to $E_F=0$ meaning that the HOMO level is the smallest negative eigenvalue, while the LUMO is the smallest positive eigenvalue:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "es = H.eigenstate()\n",
    "\n",
    "# We specify an origin to center the molecule in the grid\n",
    "h2o.lattice.origin = [-4, -4, -4]\n",
    "\n",
    "# Reduce the contained eigenstates to only the HOMO and LUMO\n",
    "# Find the index of the smallest positive eigenvalue\n",
    "idx_lumo = (es.eig > 0).nonzero()[0][0]\n",
    "es = es.sub([idx_lumo - 1, idx_lumo])\n",
    "\n",
    "plots = [\n",
    "    h2o.plot(axes=\"xy\", atoms_style={\"size\": normalize(n), \"color\": c})\n",
    "    for n, c in zip(\n",
    "        h2o.apply(\n",
    "            es.norm2(projection=\"orbital\"),\n",
    "            np.sum,\n",
    "            mapper=partial(h2o.a2o, all=True),\n",
    "            axis=1,\n",
    "        ),\n",
    "        (\"red\", \"blue\", \"green\"),\n",
    "    )\n",
    "]\n",
    "\n",
    "merge_plots(*plots, composite_method=\"subplots\", cols=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "These are not that interesting. The projection of the HOMO and LUMO states show where the largest weight of the HOMO and LUMO states, however we can't see the orbital symmetry differences between the HOMO and LUMO states.\n",
    "\n",
    "Instead of plotting the weight on each orbital it is more interesting to plot the actual wavefunctions which contains the orbital symmetries, however, matplotlib is currently not capable of plotting real-space iso-surface plots. To do this, please use VMD or your preferred software."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def integrate(g):\n",
    "    print(\n",
    "        \"Real space integrated wavefunction: {:.4f}\".format(\n",
    "            (np.absolute(g.grid) ** 2).sum() * g.dvolume\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "g = si.Grid(0.2, lattice=h2o.lattice)\n",
    "es.sub(0).wavefunction(g)\n",
    "integrate(g)\n",
    "# g.write('HOMO.cube')\n",
    "g.fill(0)  # reset the grid values to 0\n",
    "es.sub(1).wavefunction(g)\n",
    "integrate(g)\n",
    "# g.write('LUMO.cube')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Real space charge\n",
    "\n",
    "Since we have the basis functions we can also plot the charge in the grid. We can do this via either reading the density matrix or read in the charge output directly from Siesta.\n",
    "Since both should yield the same value we can compare the output from Siesta with that calculated in `sisl`.\n",
    "\n",
    "You will notice that re-creating the density on a real space grid in sisl is much slower than creating the wavefunction. This is because we need _orbital multiplications_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "DM = fdf.read_density_matrix()\n",
    "rho = si.get_sile(\"siesta_1.nc\").read_grid(\"Rho\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a new grid that is the negative of the readed charge density, then use sisl to calculate the charge density. This will check whether sisl and Siesta creates the same charge density."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "diff = rho * (-1)\n",
    "DM.density(diff)\n",
    "print(\n",
    "    \"Real space integrated density difference: {:.3e}\".format(\n",
    "        diff.grid.sum() * diff.dvolume\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.11.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}