{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "os.chdir(\"siesta_2\")\n",
    "import numpy as np\n",
    "from sisl import *\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Siesta ---  graphene\n",
    "\n",
    "This tutorial will describe a complete walk-through of a large fraction of the `sisl` functionalities that may be related to the [Siesta code](https://gitlab.com/siesta-project/siesta).\n",
    "\n",
    "Contrary to the $\\mathrm H_2\\mathrm O$ system this tutorial will emphasize the usefulness of performing bandstructures etc. directly in Python using sisl.\n",
    "\n",
    "## Creating the geometry\n",
    "\n",
    "Our system of interest will be the smallest graphene cell. Instead of defining the atomic positions, the carbon atoms and supercell for graphene, we use a default implementation of graphene in `sisl`. There are a small selection of the typical geometries, including graphene."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graphene = geom.graphene(1.44)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us plot the geometry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "graphene.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_2\n",
    "PAO.BasisSize SZP\n",
    "MeshCutoff 250. Ry\n",
    "CDF.Save true\n",
    "CDF.Compress 9\n",
    "SaveHS true\n",
    "SaveRho true\n",
    "%block kgrid.MonkhorstPack\n",
    "  61  1 1 0.\n",
    "   1 61 1 0.\n",
    "   0  0 1 0.\n",
    "%endblock\n",
    "\"\"\"\n",
    ")\n",
    "graphene.write(\"STRUCT.fdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating the electronic structure\n",
    "\n",
    "Before proceeding, run Siesta to calculate the ground state electronic structure.\n",
    "\n",
    "After having completed the Siesta run we may read the Hamiltonian to manipulate and extract different information.\n",
    "After reading the Hamiltonian it is obvious that a great deal of new data has been associated with the Hamiltonian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fdf = get_sile(\"RUN.fdf\")\n",
    "H = fdf.read_hamiltonian()\n",
    "print(H)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Calculating DOS and PDOS\n",
    "\n",
    "When we are dealing with periodic structures (such as graphene) it is imperative to calculate the density of states in a simple and efficient manner. Below we will calculate the DOS for a variety of Monkhorst-Pack grids to check the convergence of the DOS (it shouldn't take more than a minute):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "E = np.linspace(-6, 4, 500)\n",
    "for nk in [21, 41, 61, 81]:\n",
    "    bz = MonkhorstPack(H, [nk, nk, 1])\n",
    "    plt.plot(\n",
    "        E,\n",
    "        bz.apply.average.eigenvalue(wrap=lambda ev: ev.DOS(E)),\n",
    "        label=\"nk={}\".format(nk),\n",
    "    )\n",
    "plt.xlim(E[0], E[-1])\n",
    "plt.ylim(0, None)\n",
    "plt.xlabel(r\"$E - E_F$ [eV]\")\n",
    "plt.ylabel(r\"DOS [1/eV]\")\n",
    "plt.legend();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that the k-points are converged for approximately 61 k-points using the default smearing method. The default smearing method is the Gaussian smearing technique with $\\sigma=0.1\\mathrm{eV}$. Note that intrinsically the `MonkhorstPack` grid assumes time-reversal symmetry. I.e. $\\mathbf k \\equiv -\\mathbf k$.\n",
    "\n",
    "Now we may use the Monkhorst-Pack grid for 81 points to find the projected DOS for some of the orbitals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pdos_plot = H.plot.pdos(kgrid=[81, 81, 1], data_Erange=(-6, 4), Erange=[-6, 4], nE=500)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orb_groups = [\n",
    "    {\"l\": 0, \"name\": \"s\", \"color\": \"red\"},\n",
    "    {\"l\": 1, \"m\": [-1, 1], \"name\": \"px + py\", \"color\": \"blue\"},\n",
    "    {\"l\": 1, \"m\": 0, \"name\": \"pz\", \"color\": \"green\"},\n",
    "]\n",
    "pdos_plot.update_inputs(groups=orb_groups)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As seen the $p_z$ orbitals are responsible for the DOS in a broad range of energies around the Fermi-level. This is one reason for the tight-binding models success with respect to graphene.\n",
    "\n",
    "Another way to gain information is via the so-called *fat-bands* which basically is the PDOS scaling (no broadening) on each band for the quantities we are interested in. To plot the fat-bands we need the band-structure and a projection of each state onto the requested orbitals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the band-structure\n",
    "bz = BandStructure(\n",
    "    H,\n",
    "    [[0] * 3, [2.0 / 3, 1.0 / 3, 0], [0.5, 0.5, 0], [1] * 3],\n",
    "    400,\n",
    "    names=[r\"$\\Gamma$\", r\"$K$\", r\"$M$\", r\"$\\Gamma$\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bz.plot.fatbands(groups=orb_groups, Erange=[-21, 10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- all, black\n",
    "- $p_z$, green\n",
    "- $s$, red\n",
    "- $p_x+p_y$, blue"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Hamiltonian eigenstates\n",
    "\n",
    "At this point we have plotted the $k$-averaged DOS, PDOS. We have also plotted the fat-bands (and thus the band-structure). \n",
    "\n",
    "In addition to these things we can plot the real-space eigenstates. We first plot the $\\Gamma$-point for the first 2x2 unit-cell. This $k$-point has complete unit-cell periodicity and thus the plotted wavefunction should be fully periodic along all directions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "es = H.eigenstate()\n",
    "idx_valence = (es.eig > 0).nonzero()[0][0] - 1\n",
    "# Only select the valence band state\n",
    "es = es.sub(idx_valence)\n",
    "\n",
    "# Generate a grid encompassing a 2x2 graphene unit-cell\n",
    "g = Grid(0.2, lattice=H.geometry.lattice.tile(2, 0).tile(2, 1))\n",
    "# Calculate the real-space wavefunctions\n",
    "es.wavefunction(g)\n",
    "\n",
    "# Extract the wavefunction a few Ang above the graphene plane\n",
    "# To do this we need to find the index of the corresponding z-plane.\n",
    "# The Grid.index method is useful in this regard.\n",
    "xyz = H.geometry.xyz[0, :].copy()\n",
    "xyz[2] += 1.0\n",
    "z_idx = g.index(xyz, axis=2)\n",
    "x, y = np.mgrid[: g.shape[0], : g.shape[1]]\n",
    "x, y = x * g.dcell[0, 0] + y * g.dcell[1, 0], x * g.dcell[0, 1] + y * g.dcell[1, 1]\n",
    "plt.contourf(x, y, g.grid[:, :, z_idx])\n",
    "xyz = H.geometry.tile(2, 0).tile(2, 1).xyz\n",
    "plt.scatter(xyz[:, 0], xyz[:, 1], 20, c=\"k\")\n",
    "plt.xlabel(r\"$x$ [Ang]\")\n",
    "plt.ylabel(r\"$y$ [Ang]\");"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will now try and plot the real-space wavefunction for a finite $k$. By choosing the $[1/2, 0, 0]$ point we know it must have a periodicity of 2 along the first lattice vector (this lattice vector is pointing right-up), and full periodicity along the second lattice vector. Since we have a finite $k$ the grid data-type *must* be complex because the eigenstates have complex components. And thus we will plot both the real and imaginary part."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_, axs = plt.subplots(1, 2, figsize=(16, 5))\n",
    "es = H.eigenstate([1.0 / 2, 0, 0])\n",
    "idx_valence = (es.eig > 0).nonzero()[0][0] - 1\n",
    "es = es.sub(idx_valence)\n",
    "g = Grid(0.2, dtype=np.complex128, lattice=H.geometry.lattice.tile(4, 0).tile(4, 1))\n",
    "es.wavefunction(g)\n",
    "x, y = np.mgrid[: g.shape[0], : g.shape[1]]\n",
    "x, y = x * g.dcell[0, 0] + y * g.dcell[1, 0], x * g.dcell[0, 1] + y * g.dcell[1, 1]\n",
    "axs[0].contourf(x, y, g.grid[:, :, z_idx].real)\n",
    "axs[1].contourf(x, y, g.grid[:, :, z_idx].imag)\n",
    "xyz = H.geometry.tile(4, 0).tile(4, 1).xyz\n",
    "for ax in axs:\n",
    "    ax.scatter(xyz[:, 0], xyz[:, 1], 20, c=\"k\")\n",
    "    ax.set_xlabel(r\"$x$ [Ang]\")\n",
    "    ax.set_ylabel(r\"$y$ [Ang]\");"
   ]
  },
  {
   "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.8.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}