{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pybinding as pb\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "pb.pltutils.use_style()\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Structure-mapped data\n",
    "\n",
    "As shown in the previous section, many classes in pybinding use structure plots in a similar way. One class stands out here: [`StructureMap`](http://docs.pybinding.site/page/plotting/../_api/pybinding.StructureMap.html#pybinding.StructureMap) can be used to map any arbitrary data onto the spatial structure of a model. [`StructureMap`](http://docs.pybinding.site/page/plotting/../_api/pybinding.StructureMap.html#pybinding.StructureMap) objects are produced in two cases: as the results of various computation functions (e.g. [`Solver.calc_spatial_ldos()`](http://docs.pybinding.site/page/plotting/../_api/pybinding.solver.html#pybinding.solver.Solver.calc_spatial_ldos)) or returned from [`Model.structure_map()`](http://docs.pybinding.site/page/plotting/../_api/pybinding.Model.html#pybinding.Model.structure_map) which can map custom user data.\n",
    "\n",
    "[Download this page as a Jupyter notebook](http://docs.pybinding.site/page/_downloads/aa7fef7b872a6ae41be96da663a83e87/structuremap.ipynb)\n",
    "\n",
    "## Draw only certain hoppings\n",
    "\n",
    "Just as before, we can draw only the desired hoppings. Note that `smap` is a [`StructureMap`](http://docs.pybinding.site/page/plotting/../_api/pybinding.StructureMap.html#pybinding.StructureMap) returned by [`Solver.calc_probability()`](http://docs.pybinding.site/page/plotting/../_api/pybinding.solver.html#pybinding.solver.Solver.calc_probability)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pybinding.repository import graphene\n",
    "\n",
    "plt.figure(figsize=(7, 3))\n",
    "\n",
    "plt.subplot(121, title=\"The model\")\n",
    "model = pb.Model(graphene.monolayer(nearest_neighbors=3), graphene.hexagon_ac(1))\n",
    "model.plot(hopping={'draw_only': ['t']})\n",
    "\n",
    "plt.subplot(122, title=\"$|\\Psi|^2$\")\n",
    "solver = pb.solver.arpack(model, k=10)\n",
    "smap = solver.calc_probability(n=2)\n",
    "smap.plot(hopping={'draw_only': ['t']})\n",
    "pb.pltutils.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Slicing a structure\n",
    "\n",
    "This follows a syntax similar to numpy fancy indexing where we can give a condition as the index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 3))\n",
    "\n",
    "plt.subplot(121, title=\"Original\")\n",
    "smap.plot(hopping={'draw_only': ['t']})\n",
    "\n",
    "plt.subplot(122, title=\"Sliced: y > 0\")\n",
    "upper = smap[smap.y > 0]\n",
    "upper.plot(hopping={'draw_only': ['t']})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(7, 3))\n",
    "\n",
    "plt.subplot(121, title=\"Original: A and B\")\n",
    "smap.plot(hopping={'draw_only': ['t', 't_nn']})\n",
    "\n",
    "plt.subplot(122, title=\"Sliced: A only\")\n",
    "a_only = smap[smap.sublattices == 'A']\n",
    "a_only.plot(hopping={'draw_only': ['t', 't_nn']})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Mapping custom data\n",
    "\n",
    "The method [`Model.structure_map()`](http://docs.pybinding.site/page/plotting/../_api/pybinding.Model.html#pybinding.Model.structure_map) returns a [`StructureMap`](http://docs.pybinding.site/page/plotting/../_api/pybinding.StructureMap.html#pybinding.StructureMap) where any user-defined `data` can be mapped to the spatial positions of the lattice sites. The `data` just needs to be a 1D array with the same size as the total number of sites in the system."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.8, 3))\n",
    "\n",
    "plt.subplot(121, title=\"The model\")\n",
    "model = pb.Model(graphene.monolayer(), graphene.hexagon_ac(1))\n",
    "model.plot()\n",
    "\n",
    "plt.subplot(122, title=\"Custom color data: 2x * (y + 1)\")\n",
    "custom_data = 2 * model.system.x * (model.system.y + 1)\n",
    "smap = model.structure_map(custom_data)\n",
    "smap.plot()\n",
    "pb.pltutils.colorbar()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.8, 3))\n",
    "\n",
    "plt.subplot(121, title=\"sin(10x)\")\n",
    "smap = model.structure_map(np.sin(10 * model.system.x))\n",
    "smap.plot()\n",
    "pb.pltutils.colorbar()\n",
    "\n",
    "plt.subplot(122, title=\"cos(5y)\")\n",
    "smap = model.structure_map(np.cos(5 * model.system.y))\n",
    "smap.plot()\n",
    "pb.pltutils.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contour plots for large systems\n",
    "\n",
    "For larger systems, structure plots don’t make much sense because the details of the sites and hoppings would be too small to see. Contour plots look much better in this case."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(6.8, 3))\n",
    "model = pb.Model(graphene.monolayer(), graphene.hexagon_ac(10))\n",
    "\n",
    "plt.subplot(121, title=\"sin(x)\")\n",
    "smap = model.structure_map(np.sin(model.system.x))\n",
    "smap.plot_contourf()\n",
    "pb.pltutils.colorbar()\n",
    "\n",
    "plt.subplot(122, title=\"cos(y/2)\")\n",
    "smap = model.structure_map(np.cos(0.5 * model.system.y))\n",
    "smap.plot_contourf()\n",
    "pb.pltutils.colorbar()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Composing multiple plots\n",
    "\n",
    "Various plotting methods or even different invocations of the same method can be composed to create nice figures. For example, we may want to use different colormaps to distinguish between sublattices A and B when plotting some data on top of the structure of graphene. Below, the first pass plots only the hopping lines, the second pass draws the sites of sublattice A and the third draws sublattice B. The darkness of the color indicates the intensity of the mapped data, while blue/red distinguishes the sublattices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = pb.Model(graphene.monolayer(), graphene.hexagon_ac(1))\n",
    "custom_data = 2 * model.system.x * (model.system.y + 1)\n",
    "smap = model.structure_map(custom_data)\n",
    "\n",
    "plt.figure(figsize=(6.8, 3))\n",
    "plt.subplot(121, title=\"Regular plot\")\n",
    "smap.plot()\n",
    "\n",
    "plt.subplot(122, title=\"Composite plot\")\n",
    "smap.plot(site_radius=0)  # only draw hopping lines, no sites\n",
    "a_only = smap[smap.sublattices == \"A\"]\n",
    "a_only.plot(cmap=\"Blues\", hopping={'width': 0})  # A sites, no hoppings\n",
    "b_only = smap[smap.sublattices == \"B\"]\n",
    "b_only.plot(cmap=\"Reds\", hopping={'width': 0})  # B sites, no hoppings"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 4
}