{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "notebook-header"
    ]
   },
   "source": [
    "[![GitHub issues by-label](https://img.shields.io/github/issues-raw/pfebrer/sisl/PdosPlot?style=for-the-badge)](https://github.com/pfebrer/sisl/labels/PdosPlot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "    \n",
    "    \n",
    "PdosPlot\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": [
    "We are going to get the PDOS from a SIESTA `.PDOS` file, but we could get it from some other source, e.g. a hamiltonian."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot = sisl.get_sile(siesta_files / \"SrTiO3\" / \"unpolarized\" / \"SrTiO3.PDOS\").plot(\n",
    "    Erange=[-10, 10]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, a PDOS plot shows the total density of states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## PDOS groups\n",
    "\n",
    "There's a very important setting in the `PdosPlot`: `groups`. This setting expects a list of orbital groups, where each group is a dictionary that can specify \n",
    "- `species`\n",
    "- `atoms`\n",
    "- `orbitals` (the orbital name)\n",
    "- `n`, `l`, `m` (the quantum numbers)\n",
    "- `Z` (the Z shell of the orbital)\n",
    "- `spin`\n",
    "\n",
    "involved in the PDOS line that you want to draw. Apart from that, a group also accepts the `name`, `color`, `linewidth` and `dash` keys that manage the aesthetics of the line and `reduce`, which indicates the method to use for accumulating orbital contributions: `\"mean\"` averages over orbitals while `\"sum\"` simply accumulates all contributions. Finally, `scale` lets you multiply the DOS of the group by whatever factor you want.\n",
    "\n",
    "Here is an example of how to use the `groups` setting to create a line that displays the Oxygen 2p PDOS:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(\n",
    "    groups=[{\"name\": \"My first PDOS (Oxygen)\", \"species\": [\"O\"], \"n\": 2, \"l\": 1}]\n",
    ")\n",
    "# or (it's equivalent)\n",
    "plot.update_inputs(\n",
    "    groups=[\n",
    "        {\n",
    "            \"name\": \"My first PDOS (Oxygen)\",\n",
    "            \"species\": [\"O\"],\n",
    "            \"orbitals\": [\"2pzZ1\", \"2pzZ2\", \"2pxZ1\", \"2pxZ2\", \"2pyZ1\", \"2pyZ2\"],\n",
    "        }\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And now we are going to create three lines, one for each species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(\n",
    "    groups=[\n",
    "        {\n",
    "            \"name\": \"Oxygen\",\n",
    "            \"species\": [\"O\"],\n",
    "            \"color\": \"darkred\",\n",
    "            \"dash\": \"dash\",\n",
    "            \"reduce\": \"mean\",\n",
    "        },\n",
    "        {\n",
    "            \"name\": \"Titanium\",\n",
    "            \"species\": [\"Ti\"],\n",
    "            \"color\": \"gray\",\n",
    "            \"size\": 3,\n",
    "            \"reduce\": \"mean\",\n",
    "        },\n",
    "        {\"name\": \"Sr\", \"species\": [\"Sr\"], \"color\": \"green\", \"reduce\": \"mean\"},\n",
    "    ],\n",
    "    Erange=[-5, 5],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "It's interesting to note that the `atoms` key of each group accepts the same possibilities as the `atoms` argument of the `Geometry` methods. Therefore, **you can use indices, categories, dictionaries, strings...**\n",
    "\n",
    "For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's import the AtomZ and AtomOdd categories just to play with them\n",
    "from sisl.geom import AtomZ, AtomOdd\n",
    "\n",
    "plot.update_inputs(\n",
    "    groups=[\n",
    "        {\"atoms\": [0, 1], \"name\": \"Atoms 0 and 1\"},\n",
    "        {\"atoms\": {\"Z\": 8}, \"name\": \"Atoms with Z=8\"},\n",
    "        {\"atoms\": AtomZ(8) & ~AtomOdd(), \"name\": \"Oxygens with even indices\"},\n",
    "    ]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Easy and fast DOS splitting\n",
    "\n",
    "As you might have noticed, sometimes it might be cumbersome to build all the groups you want. If your needs are simple and you don't need the flexibility of defining every parameter by yourself, there is a set of methods that will help you explore your PDOS data faster than ever before. These are: `split_DOS`, `split_groups`, `update_groups`, `remove_groups` and `add_groups`.\n",
    "\n",
    "Let's begin with `split_DOS`. As you can imagine, this method splits the density of states:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "By default, it splits on the different species, but you can use the `on` argument to change that."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS(on=\"atoms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have the contribution of each atom.\n",
    "\n",
    "But here comes the powerful part: `split_DOS` accepts as keyword arguments all the keys that a group accepts. Then, it adds that extra constrain to the splitting by adding the value to each group. So, if we want to get the separate contributions of all oxygen atoms, **we can impose an extra constraint** on species:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS(on=\"atoms\", species=[\"O\"], name=\"Oxygen $atoms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and then we have only the oxygen atoms, which are all equivalent.\n",
    "\n",
    "Note that we also set a name for all groups, with the additional twist that we used the templating supported by `split_DOS`. If you are splitting on `parameter`, you can use `$parameter` inside your name and the method will replace it with the value for each group. In this case `parameter` was `atoms`, but it could be anything you are splitting the DOS on.\n",
    "\n",
    "You can also **exclude some values of the parameter you are splitting on**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS(on=\"atoms\", exclude=[1, 3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or **indicate the only values that you want**:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS(on=\"atoms\", only=[0, 2])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, if you want to split on multiple parameters at the same time, you can use `+` between different parameters. For example, to get all the oxygen orbitals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS(on=\"n+l+m\", species=[\"O\"], name=\"Oxygen\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Managing existing groups\n",
    "\n",
    "Not only you can create groups easily with `split_DOS`, but it's also easy to manage the groups that you have created. \n",
    "\n",
    "The methods that help you accomplish this are `split_groups`, `update_groups`, `remove_groups`. All three methods accept an undefined number of arguments that are used to select the groups you want to act on. You can refer to groups by their name (using a `str`) or their position (using an `int`). It's very easy to understand with examples. Then, keyword arguments depend on the functionality of each method.\n",
    "\n",
    "For example, let's say that we have splitted the DOS on species"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS(name=\"$species\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "and we want to remove the Sr and O lines. That's easy:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.remove_groups(\"Sr\", 2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We have indicated that we wanted to remove the group with name `\"Sr\"` and the 2nd group. Simple, isn't it?\n",
    "\n",
    "Now that we know how to indicate the groups that we want to act on, let's use it to get the total `Sr` contribution, and then the `Ti` and `O` contributions splitted by `n` and `l`.\n",
    "\n",
    "It sounds difficult, but it's actually not. Just split the DOS on species:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.split_DOS(name=\"$species\", reduce=\"mean\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And then use `split_groups` to split only the groups that we want to split:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "plot.split_groups(\"Sr\", 2, on=\"n+l\", dash=\"dot\").show(\"png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how we've also set `dash` for all the groups that `split_groups` has generated. We can do this because `split_groups` works exactly as `split_DOS`, with the only difference that splits specific groups.\n",
    "\n",
    "Just as a last thing, we will let you figure out how `update_groups` works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_groups(\"Ti\", color=\"red\", size=4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "notebook-end"
    ]
   },
   "source": [
    "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
}