{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": [
     "notebook-header"
    ]
   },
   "source": [
    "[![GitHub issues by-label](https://img.shields.io/github/issues-raw/pfebrer/sisl/AtomicmatrixPlot?style=for-the-badge)](https://github.com/pfebrer/sisl/labels/AtomicMatrixPlot)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "AtomicMatrixPlot\n",
    "=========\n",
    "\n",
    "`AtomicMatrixPlot` allows you to visualize sparse matrices. This can help you:\n",
    "\n",
    "- **Understand sparse matrices better**, if you are new to them.\n",
    "- Easily **introspect matrices** to debug or understand how to implement new functionality."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sisl\n",
    "\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's create a toy **Hamiltonian** that we will play with. We will use a chain with two atoms: \n",
    "\n",
    "  - C: With one s orbital and a set of p orbitals.\n",
    "  - H: With one s orbital."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a C atom with s and p orbitals\n",
    "C = sisl.Atom(\n",
    "    \"C\",\n",
    "    orbitals=[\n",
    "        sisl.AtomicOrbital(\"2s\", R=0.8),\n",
    "        sisl.AtomicOrbital(\"2py\", R=0.8),\n",
    "        sisl.AtomicOrbital(\"2pz\", R=0.8),\n",
    "        sisl.AtomicOrbital(\"2px\", R=0.8),\n",
    "    ],\n",
    ")\n",
    "\n",
    "# Create a H atom with one s orbital\n",
    "H = sisl.Atom(\"H\", orbitals=[sisl.AtomicOrbital(\"1s\", R=0.4)])\n",
    "\n",
    "# Create a chain along X\n",
    "geom = sisl.Geometry(\n",
    "    [[0, 0, 0], [1, 0, 0]],\n",
    "    atoms=[C, H],\n",
    "    lattice=sisl.Lattice([2, 10, 10], nsc=[3, 1, 1]),\n",
    ")\n",
    "\n",
    "# Random Hamiltonian with non-zero elements only for orbitals that overlap\n",
    "H = sisl.Hamiltonian(geom)\n",
    "for i in range(geom.no):\n",
    "    for j in range(geom.no * 3):\n",
    "        dist = geom.rij(*geom.o2a([i, j]))\n",
    "        if dist < 1.2:\n",
    "            H[i, j] = (np.random.random() - 0.5) * (1 if j < geom.no else 0.2)\n",
    "# Symmetrize it to make it more realistic\n",
    "H = H + H.transpose()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Matrix as an image\n",
    "\n",
    "The default mode of `AtomicMatrixPlot` is simply to plot an image where the values of the matrix are encoded as colors:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot = H.plot()\n",
    "plot.get()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Changing the colorscale\n",
    "\n",
    "The most obvious thing to tweak here is the colorscale, you can do so by using the `colorscale` input, which, as usual, accepts any colorscale that the plotting backend can understand."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(colorscale=\"temps\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note how the `temps` colorscale makes it more clear that there are elements of the matrix that are not set (because the matrix is sparse). Those elements are not displayed.\n",
    "\n",
    "The range of colors **by default is set from min to max**. **Unless there are negative and positive values**. In that case, the colorscale is just **centered at 0** by default.\n",
    "\n",
    "However, you can set the `crange` and `cmid` to customize the colorscale as you wish. For example, to center the scale at `0.5`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(cmid=0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "And to set the range of the scale from `-4` to `1`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(crange=(-4, 1))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Notice how **`crange` takes precedence over `cmid`**. Now, to go back to the default range, just set both to `None`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(crange=None, cmid=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Show values as text\n",
    "\n",
    "Colors are nice to give you a quick impression of the relative magnitude of the matrix elements. However, you might want to know the exact value. Although `plotly` shows them when you pass the mouse over the matrix elements, sometimes it might be more convenient to directly display them on top.\n",
    "\n",
    "To do this, you need to pass a formatting string to the `text` input. For example, to show two decimals:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(text=\".2f\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can tweak the style of text with the `textfont` input, which is a dictionary with three (optional) keys:\n",
    "\n",
    "- `color`: Text color.\n",
    "- `family`: Font family for the text. Note that different backends might support different fonts.\n",
    "- `size`: The size of the font.\n",
    "\n",
    "The default value will be used for any key that you don't include in the dictionary."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(textfont={\"color\": \"blue\", \"family\": \"times\", \"size\": 15})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want the text to go away, set the `text` input to `None`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(text=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot = plot.update_inputs(textfont={}, text=\".2f\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Separators\n",
    "\n",
    "Sparse atomic matrices may be hard to interpret because it's hard to know by eye to which atoms/orbitals an element belongs.\n",
    "\n",
    "Separators come to the rescue by providing a guide for your eye to quickly pinpoint what each element is. There are three types of separators:\n",
    "\n",
    "- `sc_lines`: Draw lines separating the **different cells of the auxiliary supercell**.\n",
    "- `atom_lines`: Draw lines separating the blocks corresponding to **each atom-atom interaction**.\n",
    "- `orbital_lines`: Within each atom, draw lines that **isolate the interactions between two sets of orbitals**. E.g. a set of 3 `p` orbitals from one atom and an `s` orbital from another atom.\n",
    "\n",
    "They all can be activated (deactivated) by setting the corresponding input to `True` (`False`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(\n",
    "    sc_lines=True,\n",
    "    atom_lines=True,\n",
    "    orbital_lines=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sometimes, the **default styles for the lines might not suit your visualization**. For example, they might not play well with your chosen colorscale. In that case, you can **pass a dictionary of line styles** to the inputs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(\n",
    "    orbital_lines={\"color\": \"pink\", \"width\": 5, \"dash\": \"dash\", \"opacity\": 0.8},\n",
    "    sc_lines={\"width\": 4},\n",
    "    atom_lines={\"color\": \"gray\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Labels\n",
    "\n",
    "You might want to have a clearer idea of the orbitals that correspond to a given matrix element. You can turn on labels with `set_labels`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(set_labels=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Labels have the format: `Atom index: (l, m)`. where l and m are the quantum numbers of the orbital."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot = plot.update_inputs(set_labels=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Showing only one cell\n",
    "\n",
    "If you only want to visualize a given cell in the supercell, you can pass the index to the `isc` input."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(isc=1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To go back to visualizing the whole supercell, just set `isc` to `None`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(isc=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Arrows\n",
    "\n",
    "One can ask for an arrow to be drawn on each matrix element. You are free to represent whatever you like as arrows.\n",
    "\n",
    "The arrow specification works the same as for atom arrows in `GeometryPlot`. It is a dictionary with the key `data` containing the arrow data and the styling keys (`color`, `width`, `opacity`...) to tweak the style.\n",
    "\n",
    "However, there is one main difference. If `data` is skipped, vertical arrows are drawn, with the value of the corresponding matrix element defining the length of the arrow:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "nbsphinx-thumbnail"
    ]
   },
   "outputs": [],
   "source": [
    "plot.update_inputs(arrows={\"color\": \"blue\"}).show(\"png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Arrows are normalized so that they fit the box** of their matrix element.\n",
    "\n",
    "It may be that pixel colors and numbers make it difficult to visualize the arrows. In that case, you can disable them both. We have already seen how to disable text. For pixel colors there's the `color_pixels` input, which is a switch to turn them on or off:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(color_pixels=False, text=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you want to plot custom data for the arrows, you have to pass a sparse matrix where the last dimension is the cartesian coordinate `(X, Y)`. Let's create some random data to display it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a sparse matrix with the same sparsity pattern as our Hamiltonian,\n",
    "# but with an extra dimension that will host the X and Y coordinates.\n",
    "arrow_data = sisl.SparseCSR.fromsp([H, H])\n",
    "\n",
    "# The X coordinate will be the Hamiltonian's value,\n",
    "# while the Y coordinate will be just random.\n",
    "for i in range(arrow_data.shape[0]):\n",
    "    for j in range(arrow_data.shape[1]):\n",
    "        if arrow_data[i, j, 1] != 0:\n",
    "            arrow_data[i, j, 1] *= np.random.random()\n",
    "\n",
    "# Let's display the data\n",
    "plot.update_inputs(arrows={\"data\": arrow_data, \"color\": \"red\"})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, we showcase how **you can have multiple specifications of arrows**.\n",
    "\n",
    "We also show how you can use the `center` key. You can set `center` to `\"start\"`, `\"middle\"` or `\"end\"`. It determines which part of the arrow is pinned to the center of the matrix element (the default is `\"middle\"`):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plot.update_inputs(\n",
    "    arrows=[\n",
    "        {\"color\": \"blue\", \"center\": \"start\", \"name\": \"Hamiltonian value\"},\n",
    "        {\"data\": arrow_data, \"color\": \"red\", \"center\": \"end\", \"name\": \"Some data\"},\n",
    "    ]\n",
    ")"
   ]
  }
 ],
 "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
}