{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "notebook-header"
    ]
   },
   "source": [
    "[![GitHub issues by-label](https://img.shields.io/github/issues-raw/pfebrer/sisl/BandsPlot?style=for-the-badge)](https://github.com/pfebrer/sisl/labels/BandsPlot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    \n",
    "    \n",
    "BandsPlot \n",
    "========="
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sisl\n",
    "\n",
    "# This is just for convenience to retreive files\n",
    "siesta_files = sisl._environ.get_environ_variable(\"SISL_FILES_TESTS\") / \"siesta\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's get a bands_plot from a `.bands` file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot = sisl.get_sile(\n",
    "    siesta_files / \"SrTiO3\" / \"unpolarized\" / \"SrTiO3.bands\"\n",
    ").plot()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and see what we've got:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Getting the bands that you want\n",
    "\n",
    "By default, `BandsPlot` gives you **the 15 bands below and above 0 eV** (which is interpreted as the fermi level). \n",
    "\n",
    "There are two main ways to specify the bands that you want to display: `Erange` and `bands_range`.\n",
    "\n",
    "As you may have guessed, `Erange` specifies the energy range that is displayed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(Erange=[-10, 10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "while with `bands_range` you can actually indicate the indices.\n",
    "\n",
    "However, note that **`Erange` has preference over `bands_range`**, therefore you need to set `Erange` to `None` if you want the change to take effect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(bands_range=[6, 15], Erange=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If your fermi level is not correctly set or you want a different energy reference, you can provide a value for `E0` to specify where your 0 should be and the bands to display will be automatically calculated from that. \n",
    "\n",
    "However, if you want to update `E0` after the plot has been build and you want `BandsPlot` to recalculate the bands for you you will need to set `Erange` and `bands_range` to `None` again."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(E0=-10, bands_range=None, Erange=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how only 25 bands are displayed now: the only 10 that are below 0 eV (there are no lower states) and 15 above 0 eV."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set them back to \"normal\"\n",
    "bands_plot = bands_plot.update_inputs(E0=0, bands_range=None, Erange=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice that in spin polarized bands, **you can select the spins to display using the `spin` setting**, just pass a list of spin components (e.g. `spin=[0]`)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Individual bands in legend\n",
    "--------------------------\n",
    "\n",
    "Usually, showing all bands individually in the legend would be too messy. However, you might want to do it so that you can interactively hide show certain bands. If that is the case, you can set `group_legend` to `False`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(group_legend=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot = bands_plot.update_inputs(group_legend=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Bands styling\n",
    "\n",
    "If all you want is to change the color and width of the bands, there's one simple solution: use the `bands_style` input to tweak the line styles.\n",
    "\n",
    "Let's show them in red:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(bands_style={\"color\": \"red\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now in green but also make them wider:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(bands_style={\"color\": \"green\", \"width\": 3})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you have spin polarized bands, `bands_style` will tweak the colors for the first spin channel, while the second one can be tuned with `spindown_style`.\n",
    "\n",
    "Finally, you can pass functions to the keys of `bands_style` to customize the styles on a band basis, or even on a point basis. The functions should accept `data` as an argument, which will be an `xarray.Dataset` containing all the bands data. It should then return a single value or an array of values. It is best shown with examples. Let's create a function just to see what we receive as an input:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def color(data):\n",
    "    \"\"\"Dummy function to see what we receive.\"\"\"\n",
    "    print(data)\n",
    "    return \"green\"\n",
    "\n",
    "\n",
    "bands_plot.update_inputs(bands_style={\"color\": color})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So, you can see that we receive a `Dataset`. The most important variable is `E`, which contains the energy (that depends on `k` and `band`). Let's now play with it to do some custom styling:\n",
    "- The **color** will be determined by **the slope of the band**.\n",
    "- We will plot **bands that are closer to the fermi level bigger** because they are more important."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def gradient(data):\n",
    "    \"\"\"Function that computes the absolute value of dE/dk.\n",
    "\n",
    "    This returns a two dimensional array (gradient depends on k and band)\n",
    "    \"\"\"\n",
    "    return abs(data.E.differentiate(\"k\"))\n",
    "\n",
    "\n",
    "def band_closeness_to_Ef(data):\n",
    "    \"\"\"Computes how close one band is to the fermi level.\n",
    "\n",
    "    This returns a one dimensional array (distance depends only on band)\n",
    "    \"\"\"\n",
    "    dist_from_Ef = abs(data.E).min(\"k\")\n",
    "\n",
    "    return (1 / dist_from_Ef**0.4) * 5\n",
    "\n",
    "\n",
    "# Now we are going to set the width of the band according to the distance from the fermi level\n",
    "# and the color according to the gradient. We are going to set the colorscale also, instead of using\n",
    "# the default one.\n",
    "bands_plot.update_inputs(\n",
    "    bands_style={\"width\": band_closeness_to_Ef, \"color\": gradient},\n",
    "    colorscale=\"temps\",\n",
    "    Erange=[-10, 10],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that by providing callables the possibilities are endless, you are only limited by your imagination!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot = bands_plot.update_inputs(bands_style={})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Displaying the smallest gaps\n",
    "\n",
    "The easiest thing to do is to let `BandsPlot` discover where the (minimum) gaps are.\n",
    "\n",
    "This is indicated by setting the `gap` parameter to `True`. One can also use `gap_color` if a particular color is desired."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(gap=True, gap_color=\"green\", Erange=[-10, 10])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This displays the minimum gaps. However there may be some issues with it: it will show **all** gaps with the minimum value. That is, if you have repeated points in the brillouin zone it will display multiple gaps that are equivalent. \n",
    "\n",
    "What's worse, if the region where your gap is is very flat, two consecutive points might have the same energy. Multiple gaps will be displayed one glued to another.\n",
    "\n",
    "To help cope with this issues, you have the `direct_gaps_only` and `gap_tol`.\n",
    "\n",
    "In this case, since we have no direct gaps, setting `direct_gaps_only` will hide them all:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(direct_gaps_only=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This example is not meaningful for `gap_tol`, but it is illustrative of what `gap_tol` does. It is the **minimum k-distance between two points to consider them \"the same point\"** in the sense that only one of them will be used to show the gap. In this case, if we set `gap_tol` all the way up to 3, the plot will consider the two gamma points to be part of the same \"point\" and therefore it will only show the gap once."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(direct_gaps_only=False, gap_tol=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is not what `gap_tol` is meant for, since it is thought to remediate the effect of locally flat bands, but still you can get the idea of what it does."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot = bands_plot.update_inputs(gap=False, gap_tol=0.01)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Displaying custom gaps\n",
    "\n",
    "If you are not happy with the gaps that the plot is displaying for you or **you simply want gaps that are not the smallest ones**, you can always use `custom_gaps`.\n",
    "\n",
    "Custom gaps should be a list where each item specifies how to draw that given gap. The key labels of each item are `from` and `to`, which specifies the k-points through which you want to draw the gap. The rest of labels are the typical styling labels: `color`, `width`...\n",
    "\n",
    "For example, if we want to plot the gamma-gamma gap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(custom_gaps=[{\"from\": \"Gamma\", \"to\": \"Gamma\", \"color\": \"red\"}])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how we got the gap probably not where we wanted, since it would be better to have it in the middle `Gamma` point, which is more visible. Instead of the K point name, you can also pass the K value.\n",
    "\n",
    "Now, you'll be happy to know that you can easily access the k values of all labels, as they are stored as part of the attributes of the `k` coordinate in the bands dataarray:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.nodes[\"bands_data\"].get().k.axis"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now all we need to do is to grab the value for the second gamma point:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "axis_info = bands_plot.nodes[\"bands_data\"].get().k.axis\n",
    "\n",
    "gap_k = None\n",
    "for val, label in zip(axis_info[\"tickvals\"], axis_info[\"ticktext\"]):\n",
    "    if label == \"Gamma\":\n",
    "        gap_k = val\n",
    "gap_k"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And use it to build a custom gap:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bands_plot.update_inputs(custom_gaps=[{\"from\": gap_k, \"to\": gap_k, \"color\": \"orange\"}])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Displaying spin texture\n",
    "\n",
    "If your bands plot comes from a non-colinear spin calculation (or is using a `Hamiltonian` with non-colinear spin), you can pass `\"x\"`, `\"y\"` or `\"z\"` to the `spin` setting in order to get a display of the spin texture.\n",
    "\n",
    "Let's read in a hamiltonian coming from a spin orbit SIESTA calculation, which is obtained from [this fantastic spin texture tutorial](https://github.com/juijan/TopoToolsSiesta/tree/master/Tutorials/Exercise/TI_02):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sisl\n",
    "\n",
    "siesta_files = sisl._environ.get_environ_variable(\"SISL_FILES_TESTS\") / \"siesta\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "H = sisl.get_sile(siesta_files / \"Bi_hexagonal\" / \"Bi_hexagonal.fdf\").read_hamiltonian()\n",
    "H.spin.is_spinorbit"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Generate the path for our band structure:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "band_struct = sisl.BandStructure(\n",
    "    H,\n",
    "    points=[\n",
    "        [1.0 / 2, 0.0, 0.0],\n",
    "        [0.0, 0.0, 0.0],\n",
    "        [1.0 / 3, 1.0 / 3, 0.0],\n",
    "        [1.0 / 2, 0.0, 0.0],\n",
    "    ],\n",
    "    divisions=301,\n",
    "    names=[\"M\", r\"Gamma\", \"K\", \"M\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And finally generate the plot:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "spin_texture_plot = band_struct.plot.bands(Erange=[-2, 2])\n",
    "spin_texture_plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now it's time to add spin texture to these bands. Remember the section on styling bands? If you haven't checked it, take a quick look at it, because it will come handy now. The main point to take from that section for our purpose here is that each key in the styles accepts a callable.\n",
    "\n",
    "As in other cases through the `sisl.viz` module, we provide callables that will work out of the box for the most common styling. In this case, what we need is the `SpinMoment` node. We will import it and use it simply by specifying the axis."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sisl.viz.data_sources import SpinMoment\n",
    "\n",
    "spin_texture_plot.update_inputs(bands_style={\"color\": SpinMoment(\"x\"), \"width\": 3})\n",
    "\n",
    "# We hide the legend so that the colorbar can be easily seen.\n",
    "spin_texture_plot.update_layout(showlegend=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "There is nothing magic about the `SpinMoment` node. If you pass a dummy callable as we did in the styling section, you will see that the bands data now contains a `spin_moments` variable since it comes from a non-colinear calculation. It is just a matter of grabbing that variable:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def color(data):\n",
    "    \"\"\"Dummy function to see what we receive.\"\"\"\n",
    "    print(data)\n",
    "    return \"green\"\n",
    "\n",
    "\n",
    "spin_texture_plot.update_inputs(bands_style={\"color\": color})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that, as shown in the styling section, you can use the `colorscale` input to change the colorscale, or use the `SpinMoment` node for the other styling keys. For example, we can set the width of the band to display whether there is some spin moment, and the color can show the sign."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "spin_texture_plot.update_inputs(\n",
    "    bands_style={\"color\": SpinMoment(\"x\"), \"width\": abs(SpinMoment(\"x\")) * 40}\n",
    ").update_layout(showlegend=False).show(\"png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "notebook-end"
    ]
   },
   "source": [
    "Notice how we did some postprocessing to adapt the values of the spin moment to some number that is suitable for the width. This is possible thanks to the magic of nodes!\n",
    "\n",
    "We hope you enjoyed what you learned!"
   ]
  }
 ],
 "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
}