matplotlib-demo.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Cottrell Visualization Demo - matplotlib\n",
    "\n",
    "This demo shows some visualization examples for the python library [matplotlib](https://matplotlib.org/)  using Hydrogen atom orbitals as example data. For each library, 1D line plots, 2D contour plots, and 3D surface plots are considered."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import itertools\n",
    "\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from wavefunctions import spherical_to_cartesian, cartesian_to_spherical, wavefunction"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Making Plots with matplotlib\n",
    "\n",
    "matplotlib is one of the oldest and most popular python plotting libraries (started 2003). Can be used with an interface that uses figure objects, or with a \"procedural\" interface which is made to mimic Matlab's plotting interface."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Generate sample data for line plot section\n",
    "\n",
    "r = np.linspace(0, 15, 30)\n",
    "\n",
    "s_wavefunctions = []\n",
    "for n in range(1, 4):\n",
    "    s_wavefunctions.append(wavefunction(n, 0, 0)(r))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demonstration of line plots with matplotlib\n",
    "\n",
    "# You can make a matplotlib graph more interactive by using %matplotlib notebook. This also adds a label over\n",
    "# the figure (Figure 1, Figure 2, etc.)\n",
    "%matplotlib notebook\n",
    "\n",
    "# Procedural interface - each command adds to the active figure.\n",
    "plt.figure() # Creates a new figure.\n",
    "plt.plot(r, s_wavefunctions[0], label='1s')\n",
    "plt.plot(r, s_wavefunctions[1], label='2s')\n",
    "plt.plot(r, s_wavefunctions[2], label='3s')\n",
    "plt.xlim(0, 10)\n",
    "plt.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can make the same plot using figure objects\n",
    "fig = plt.figure() # Create a new figure\n",
    "ax = fig.add_subplot(111)\n",
    "ax.plot(r, s_wavefunctions[0], label='1s')\n",
    "ax.plot(r, s_wavefunctions[1], label='2s')\n",
    "ax.plot(r, s_wavefunctions[2], label='3s')\n",
    "ax.set_xlim(0, 10)\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# These seem pretty similar, but the figure object interface comes in handy when working with subplots.\n",
    "\n",
    "# We can use the figure object interface to create a plot with multiple subplots. \n",
    "fig, ax = plt.subplots(1,3) # First argument is the number of rows, seconds is number of columns.\n",
    "\n",
    "# ax is a list of the axes, you can add data to them the same way you did above\n",
    "ax[0].plot(r, s_wavefunctions[0], label='1s')\n",
    "ax[1].plot(r, s_wavefunctions[1], label='2s')\n",
    "ax[2].plot(r, s_wavefunctions[2], label='3s')\n",
    "\n",
    "# Use tight_layout to avoid overlapping labels\n",
    "plt.tight_layout()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Contour Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make grid on xy plane for contour plot\n",
    "side = np.linspace(-20, 20, 41)\n",
    "num_points_side = len(side)\n",
    "combinations = np.array(list(itertools.product(side, side, [0])))\n",
    "x, y, z = combinations[:,0], combinations[:,1], combinations[:,2]\n",
    "r, theta, phi = cartesian_to_spherical(x, y, z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get px, py, pz on xy plane\n",
    "p_wavefunctions_xy = []\n",
    "for i in range(-1, 2):\n",
    "    p_wavefunction = wavefunction(2, 1, i)(r, theta, phi).reshape(num_points_side, num_points_side)\n",
    "    p_wavefunctions_xy.append(p_wavefunction)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2p orbitals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contour_min = -p_wavefunctions_xy[0].max()\n",
    "contour_max= p_wavefunctions_xy[0].max()\n",
    "\n",
    "# Create subplots \n",
    "fig, ax = plt.subplots(1,3, figsize=(9,4)) # Add figsize argument here to set figure side\n",
    "\n",
    "# ax is a list of the axes, you can add data to them the same way you did above.\n",
    "# we add a color map specification here (Red to Blue), and minimum and maximum color map values.\n",
    "for i in range(3):\n",
    "    CS = ax[i].contour(side, side, p_wavefunctions_xy[i], cmap='RdBu', vmin=contour_min, vmax=contour_max)\n",
    "    ax[i].set_aspect('equal')\n",
    "\n",
    "\n",
    "fig.colorbar(CS, ax=ax[:], location='bottom')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can also make the orbitals filled\n",
    "\n",
    "contour_min = -p_wavefunctions_xy[0].max()\n",
    "contour_max= p_wavefunctions_xy[0].max()\n",
    "\n",
    "# Create subplots \n",
    "fig, ax = plt.subplots(1,3, figsize=(10,4)) # Add figsize argument here to set figure side\n",
    "\n",
    "# ax is a list of the axes, you can add data to them the same way you did above.\n",
    "# we add a color map specification here (Red to Blue), and minimum and maximum color map values.\n",
    "for i in range(3):\n",
    "    CS = ax[i].contourf(side, side, p_wavefunctions_xy[i], cmap='RdBu', vmin=contour_min, vmax=contour_max)\n",
    "    ax[i].set_aspect('equal')\n",
    "\n",
    "fig.colorbar(CS, ax=ax[:], location='bottom')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3p orbitals"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Get px, py, pz on xy plane\n",
    "p_wavefunctions_xy = []\n",
    "for i in range(-1, 2):\n",
    "    p_wavefunction = wavefunction(3, 1, i)(r, theta, phi).reshape(num_points_side, num_points_side)\n",
    "    p_wavefunctions_xy.append(p_wavefunction)\n",
    "    \n",
    "contour_min = -p_wavefunctions_xy[0].max()\n",
    "contour_max= p_wavefunctions_xy[0].max()\n",
    "\n",
    "# Create subplots \n",
    "fig, ax = plt.subplots(1,3, figsize=(10,4)) # Add figsize argument here to set figure side\n",
    "\n",
    "# ax is a list of the axes, you can add data to them the same way you did above.\n",
    "# we add a color map specification here (Red to Blue), and minimum and maximum color map values.\n",
    "for i in range(3):\n",
    "    CS = ax[i].contourf(side, side, p_wavefunctions_xy[i], cmap='RdBu', vmin=contour_min, vmax=contour_max)\n",
    "    ax[i].set_aspect('equal')\n",
    "\n",
    "fig.colorbar(CS, ax=ax[:], location='bottom')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This isn't quite what we expect. Looking at the figure, we see that this may be because our contour plot doesn't have enough levels. We can use linspace to create an evenly spaced list of levels, then add the `levels` argument to the `contourf` function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "contour_min = -p_wavefunctions_xy[0].max()\n",
    "contour_max= p_wavefunctions_xy[0].max()\n",
    "levels = np.linspace(contour_min, contour_max, 30)\n",
    "\n",
    "# Get px, py, pz on xy plane\n",
    "p_wavefunctions_xy = []\n",
    "for i in range(-1, 2):\n",
    "    p_wavefunction = wavefunction(3, 1, i)(r, theta, phi).reshape(num_points_side, num_points_side)\n",
    "    p_wavefunctions_xy.append(p_wavefunction)\n",
    "\n",
    "#contour_min = -p_wavefunctions_xy[0].max()\n",
    "#contour_max= p_wavefunctions_xy[0].max()\n",
    "\n",
    "# Create subplots \n",
    "fig, ax = plt.subplots(1,3, figsize=(10,4)) # Add figsize argument here to set figure side\n",
    "\n",
    "# ax is a list of the axes, you can add data to them the same way you did above.\n",
    "# we add a color map specification here (Red to Blue), and minimum and maximum color map values.\n",
    "for i in range(3):\n",
    "    CS = ax[i].contourf(side, side, p_wavefunctions_xy[i], cmap='RdBu', vmin=contour_min, vmax=contour_max, levels=levels)\n",
    "    ax[i].set_aspect('equal')\n",
    "\n",
    "fig.colorbar(CS, ax=ax[:], location='bottom')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Surface Plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Make grid in 3D space\n",
    "side = np.linspace(-20, 20, 41)\n",
    "num_points_side = len(side)\n",
    "combinations = np.array(list(itertools.product(side, side, side)))\n",
    "x, y, z = combinations[:,0], combinations[:,1], combinations[:,2]\n",
    "r, theta, phi = cartesian_to_spherical(x, y, z)\n",
    "\n",
    "# 3p wavefunction\n",
    "p3 = wavefunction(3, 1, 0)(r, theta, phi)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Matplotlib does not have an isosurface plot, so one option might be to make a 3D scatter plot. One downside of a 3D scatter plot in matplotlib is that you cannot set the axis ratios as with the contour plots (`axis equal`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mpl_toolkits.mplot3d import Axes3D\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "\n",
    "# Filter data \n",
    "mask_low = p3 < np.percentile(p3, 1)\n",
    "mask_high = p3 > np.percentile(p3, 99) \n",
    "x_filtered = x[mask_low | mask_high]\n",
    "y_filtered = y[mask_low | mask_high]\n",
    "z_filtered = z[mask_low | mask_high]\n",
    "p3_filtered = p3[mask_low | mask_high]\n",
    "\n",
    "ax.scatter(x_filtered, y_filtered, z_filtered, c=p3_filtered, cmap='RdBu', )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}