ch04-code-listing.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 4: Plotting and visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Robert Johansson\n",
    "\n",
    "Source code listings for [Numerical Python - A Practical Techniques Approach for Industry](http://www.apress.com/9781484205549) (ISBN 978-1-484205-54-9).\n",
    "\n",
    "The source code listings can be downloaded from http://www.apress.com/9781484205549"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %config InlineBackend.figure_format='retina'\n",
    "# %config InlineBackend.figure_format='svg'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gettings started"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-5, 2, 100)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y1 = x**3 + 5*x**2 + 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y2 = 3*x**2 + 10*x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "y3 = 6*x + 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(x, y1, color=\"blue\", label=\"y(x)\")\n",
    "ax.plot(x, y2, color=\"red\", label=\"y'(x)\")\n",
    "ax.plot(x, y3, color=\"green\", label=\"y''(x)\")\n",
    "\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.legend()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig.savefig(\"ch4-figure-1.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mpl.rcParams[\"font.family\"] = \"serif\"\n",
    "mpl.rcParams[\"font.size\"] = \"12\"\n",
    "\n",
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(x, y1, lw=1.5, color=\"blue\", label=r\"$y(x)$\")\n",
    "ax.plot(x, y2, lw=1.5, color=\"red\", label=r\"$y'(x)$\")\n",
    "ax.plot(x, y3, lw=1.5, color=\"green\", label=r\"$y''(x)$\")\n",
    "\n",
    "ax.plot(x, np.zeros_like(x), lw=0.5, color=\"black\")\n",
    "ax.plot([-3.33, -3.33], [0, (-3.3)**3 + 5*(-3.3)**2 + 10], ls='--', lw=0.5, color=\"black\")\n",
    "ax.plot([0, 0], [0, 10], lw=0.5, ls='--', color=\"black\")\n",
    "ax.plot([0], [10], lw=0.5, marker='o', color=\"blue\")\n",
    "ax.plot([-3.33], [(-3.3)**3 + 5*(-3.3)**2 + 10], lw=0.5, marker='o', color=\"blue\")\n",
    "\n",
    "ax.set_ylim(-15, 40)\n",
    "ax.set_yticks([-10, 0, 10, 20, 30])\n",
    "ax.set_xticks([-4, -2, 0, 2])\n",
    "\n",
    "ax.set_xlabel(\"$x$\", fontsize=18)\n",
    "ax.set_ylabel(\"$y$\", fontsize=18)\n",
    "ax.legend(loc=0, ncol=3, fontsize=14, frameon=False)\n",
    "\n",
    "fig.tight_layout();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig.savefig(\"ch4-figure-2.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "mpl.rcParams[\"font.family\"] = \"sans-serif\"\n",
    "mpl.rcParams[\"font.size\"] = \"10\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Backends"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "#%config InlineBackend.figure_format='svg'\n",
    "%config InlineBackend.figure_format='retina'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import matplotlib as mpl\n",
    "#mpl.use('qt4agg')\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-5, 2, 100)\n",
    "y1 = x**3 + 5*x**2 + 10\n",
    "y2 = 3*x**2 + 10*x\n",
    "y3 = 6*x + 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "ax.plot(x, y1, color=\"blue\", label=\"y(x)\")\n",
    "ax.plot(x, y2, color=\"red\", label=\"y'(x)\")\n",
    "ax.plot(x, y3, color=\"green\", label=\"y''(x)\")\n",
    "\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "ax.legend()\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Figure"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 2.5), facecolor=\"#f1f1f1\")\n",
    "\n",
    "# axes coordinates as fractions of the canvas width and height\n",
    "left, bottom, width, height = 0.1, 0.1, 0.8, 0.8\n",
    "ax = fig.add_axes((left, bottom, width, height), facecolor=\"#e1e1e1\")\n",
    "\n",
    "x = np.linspace(-2, 2, 1000)\n",
    "y1 = np.cos(40 * x)\n",
    "y2 = np.exp(-x**2)\n",
    "\n",
    "ax.plot(x, y1 * y2)\n",
    "ax.plot(x, y2, 'g')\n",
    "ax.plot(x, -y2, 'g')\n",
    "ax.set_xlabel(\"x\")\n",
    "ax.set_ylabel(\"y\")\n",
    "\n",
    "fig.savefig(\"graph.png\", dpi=100, facecolor=\"#f1f1f1\")\n",
    "fig.savefig(\"graph.pdf\", dpi=300, facecolor=\"#f1f1f1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot types"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fignum = 0\n",
    "\n",
    "def hide_labels(fig, ax):\n",
    "    global fignum\n",
    "    ax.set_xticks([])\n",
    "    ax.set_yticks([])\n",
    "    ax.xaxis.set_ticks_position('none')\n",
    "    ax.yaxis.set_ticks_position('none')\n",
    "    ax.axis('tight')\n",
    "    \n",
    "    fignum += 1\n",
    "    fig.savefig(\"plot-types-%d.pdf\" % fignum)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-3, 3, 25)\n",
    "y1 = x**3+ 3 * x**2 + 10\n",
    "y2 = -1.5 * x**3 + 10*x**2 - 15"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "\n",
    "ax.plot(x, y1)\n",
    "ax.plot(x, y2)\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "\n",
    "ax.step(x, y1)\n",
    "ax.step(x, y2)\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "width = 6/50.0\n",
    "ax.bar(x - width/2, y1, width=width, color=\"blue\")\n",
    "ax.bar(x + width/2, y2, width=width, color=\"green\")\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "ax.fill_between(x, y1, y2)\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "ax.hist(y2, bins=30)\n",
    "ax.hist(y1, bins=30)\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "\n",
    "ax.errorbar(x, y2, yerr=y1, fmt='o-')\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "\n",
    "ax.stem(x, y2, 'b', markerfmt='bs')\n",
    "ax.stem(x, y1, 'r', markerfmt='ro')\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(4, 3))\n",
    "\n",
    "x = np.linspace(0, 5, 50)\n",
    "\n",
    "ax.scatter(x, -1 + x + 0.25 * x**2 + 2 * np.random.rand(len(x)))\n",
    "ax.scatter(x, np.sqrt(x) + 2 * np.random.rand(len(x)), color=\"green\")\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(3, 3))\n",
    "\n",
    "colors = ['yellowgreen', 'gold', 'lightskyblue', 'lightcoral']\n",
    "\n",
    "x = y = np.linspace(-2, 2, 10)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "U = np.sin(X)\n",
    "V = np.sin(Y)\n",
    "\n",
    "ax.quiver(X, Y, U, V)\n",
    "\n",
    "hide_labels(fig, ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Text formatting and annotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "x = np.linspace(-20, 20, 100)\n",
    "y = np.sin(x) / x\n",
    "\n",
    "ax.plot(x, y)\n",
    "\n",
    "ax.set_ylabel(\"y label\")\n",
    "ax.set_xlabel(\"x label\")\n",
    "\n",
    "for label in ax.get_xticklabels() + ax.get_yticklabels():\n",
    "    label.set_rotation(45)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(12, 3))\n",
    "\n",
    "ax.set_yticks([])\n",
    "ax.set_xticks([])\n",
    "ax.set_xlim(-0.5, 3.5)\n",
    "ax.set_ylim(-0.05, 0.25)\n",
    "ax.axhline(0)\n",
    "\n",
    "ax.text(0, 0.1, \"Text label\", fontsize=14, family=\"serif\")\n",
    "\n",
    "ax.plot(1, 0, 'o')\n",
    "ax.annotate(\"Annotation\",\n",
    "            fontsize=14, family=\"serif\",\n",
    "            xy=(1, 0), xycoords='data',\n",
    "            xytext=(+20, +50), textcoords='offset points', \n",
    "            arrowprops=dict(arrowstyle=\"->\", connectionstyle=\"arc3, rad=.5\"))\n",
    "\n",
    "ax.text(2, 0.1, r\"Equation: $i\\hbar\\partial_t \\Psi = \\hat{H}\\Psi$\", fontsize=14, family=\"serif\")\n",
    "\n",
    "fig.savefig(\"ch4-text-annotation.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Axes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(ncols=2, nrows=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Line properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "list(range(1, 12, 2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "s.sin(sym_x).series(n=4).removeO()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "import sympy as s\n",
    "import numpy as np\n",
    "\n",
    "# a symbolic variable for x, and a numerical array with specific values of x\n",
    "sym_x = s.Symbol(\"x\")\n",
    "x = np.linspace(-2 * np.pi, 2 * np.pi, 100)\n",
    "\n",
    "def sin_expansion(x, n):\n",
    "    \"\"\"\n",
    "    Evaluate the nth order Talyor series expansion\n",
    "    of sin(x) for the numerical values in the array x.\n",
    "    \"\"\"\n",
    "    return s.lambdify(sym_x, s.sin(sym_x).series(n=n+1).removeO(), 'numpy')(x)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(10, 6))\n",
    "\n",
    "ax.plot(x, np.sin(x), linewidth=4, color=\"red\", label='sin(x)')\n",
    "\n",
    "colors = [\"blue\", \"black\"]\n",
    "linestyles = [':', '-.', '--']\n",
    "for idx, n in enumerate(range(1, 12, 2)):\n",
    "    ax.plot(x, sin_expansion(x, n), color=colors[idx // 3],\n",
    "            linestyle=linestyles[idx % 3], linewidth=3,\n",
    "            label=\"O(%d) approx.\" % (n+1))\n",
    "    \n",
    "ax.set_ylim(-1.1, 1.1)\n",
    "ax.set_xlim(-1.5*np.pi, 1.5*np.pi)\n",
    "ax.legend(bbox_to_anchor=(1.02, 1), loc=2, borderaxespad=0.0)\n",
    "fig.subplots_adjust(right=.75);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig.savefig(\"sin-expansion.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(8, 3.5), sharey=True)\n",
    "\n",
    "data1 = np.random.randn(200, 2) * np.array([3, 1])\n",
    "area1 = (np.random.randn(200) + 0.5) * 100\n",
    "\n",
    "data2 = np.random.randn(200, 2) * np.array([1, 3])\n",
    "area2 = (np.random.randn(200) + 0.5) * 100\n",
    "\n",
    "axes[0].scatter(data1[:,0], data1[:,1], color=\"green\", marker=\"s\", s=30, alpha=0.5)\n",
    "axes[0].scatter(data2[:,0], data2[:,1], color=\"blue\", marker=\"o\", s=30, alpha=0.5)\n",
    "\n",
    "axes[1].hist([data1[:,1], data2[:,1]], bins=15, color=[\"green\", \"blue\"], alpha=0.5, orientation='horizontal');"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Legends"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 4, figsize=(16, 4))\n",
    "\n",
    "x = np.linspace(0, 1, 100)\n",
    "\n",
    "for n in range(4):\n",
    "    axes[n].plot(x, x, label=\"y(x) = x\")\n",
    "    axes[n].plot(x, x + x**2, label=\"y(x) = x + x**2\")\n",
    "    axes[n].legend(loc=n+1)\n",
    "    axes[n].set_title(\"legend(loc=%d)\" % (n+1))\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"legend-loc.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(1, 1, figsize=(8.5, 3))\n",
    "\n",
    "x = np.linspace(-1, 1, 100)\n",
    "\n",
    "for n in range(1, 9):\n",
    "    ax.plot(x, n * x, label=\"y(x) = %d*x\" % n)\n",
    "\n",
    "ax.legend(ncol=4, loc=3, bbox_to_anchor=(0, 1), fontsize=12)\n",
    "fig.subplots_adjust(top=.75);\n",
    "fig.savefig(\"legend-loc-2.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Axis labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(8, 3), subplot_kw={'facecolor': \"#ebf5ff\"})\n",
    "\n",
    "x = np.linspace(0, 50, 500)\n",
    "ax.plot(x, np.sin(x) * np.exp(-x/10), lw=2)\n",
    "\n",
    "ax.set_xlabel(\"x\", labelpad=5,\n",
    "              fontsize=18, fontname='serif', color=\"blue\")\n",
    "ax.set_ylabel(\"f(x)\", labelpad=15,\n",
    "              fontsize=18, fontname='serif', color=\"blue\")\n",
    "ax.set_title(\"axis labels and title example\", loc='left',\n",
    "             fontsize=16, fontname='serif', color=\"blue\")\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-axis-labels.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Axis range"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = np.linspace(0, 30, 500)\n",
    "y = np.sin(x) * np.exp(-x/10)\n",
    "\n",
    "fig, axes = plt.subplots(1, 3, figsize=(9, 3), subplot_kw={'facecolor': \"#ebf5ff\"})\n",
    "\n",
    "axes[0].plot(x, y, lw=2)\n",
    "axes[0].set_xlim(-5, 35)\n",
    "axes[0].set_ylim(-1, 1)\n",
    "axes[0].set_title(\"set_xlim / set_y_lim\")\n",
    "\n",
    "axes[1].plot(x, y, lw=2)\n",
    "axes[1].axis('tight')\n",
    "axes[1].set_title(\"axis('tight')\")\n",
    "\n",
    "axes[2].plot(x, y, lw=2)\n",
    "axes[2].axis('equal')\n",
    "axes[2].set_title(\"axis('equal')\")\n",
    "\n",
    "fig.savefig(\"ch4-axis-ranges.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Ticks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-2 * np.pi, 2 * np.pi, 500)\n",
    "y = np.sin(x) * np.exp(-x**2/20)\n",
    "\n",
    "fig, axes = plt.subplots(1, 4, figsize=(12, 3))\n",
    "\n",
    "axes[0].plot(x, y, lw=2)\n",
    "axes[0].set_title(\"default ticks\")\n",
    "\n",
    "axes[1].plot(x, y, lw=2)\n",
    "axes[1].set_yticks([-1, 0, 1])\n",
    "axes[1].set_xticks([-5, 0, 5])\n",
    "axes[1].set_title(\"set_xticks\")\n",
    "\n",
    "axes[2].plot(x, y, lw=2)\n",
    "axes[2].xaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "axes[2].yaxis.set_major_locator(mpl.ticker.FixedLocator([-1, 0, 1]))\n",
    "axes[2].xaxis.set_minor_locator(mpl.ticker.MaxNLocator(8))\n",
    "axes[2].yaxis.set_minor_locator(mpl.ticker.MaxNLocator(8))\n",
    "axes[2].set_title(\"set_major_locator\")\n",
    "\n",
    "axes[3].plot(x, y, lw=2)\n",
    "axes[3].set_yticks([-1, 0, 1])\n",
    "axes[3].set_xticks([-2 * np.pi, -np.pi, 0, np.pi, 2 * np.pi])\n",
    "axes[3].set_xticklabels(['$-2\\pi$', '$-\\pi$', 0, r'$\\pi$', r'$2\\pi$'])\n",
    "axes[3].xaxis.set_minor_locator(mpl.ticker.FixedLocator([-3 * np.pi / 2, -np.pi/2, 0, np.pi/2, 3 * np.pi/2]))\n",
    "axes[3].yaxis.set_minor_locator(mpl.ticker.MaxNLocator(4))\n",
    "axes[3].set_title(\"set_xticklabels\")\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-axis-ticks.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(12, 4))\n",
    "\n",
    "x_major_ticker = mpl.ticker.MultipleLocator(4)\n",
    "x_minor_ticker = mpl.ticker.MultipleLocator(1)\n",
    "y_major_ticker = mpl.ticker.MultipleLocator(0.5)\n",
    "y_minor_ticker = mpl.ticker.MultipleLocator(0.25)\n",
    "\n",
    "for ax in axes:\n",
    "    ax.plot(x, y, lw=2)\n",
    "    ax.xaxis.set_major_locator(x_major_ticker)\n",
    "    ax.yaxis.set_major_locator(y_major_ticker)\n",
    "    ax.xaxis.set_minor_locator(x_minor_ticker)\n",
    "    ax.yaxis.set_minor_locator(y_minor_ticker)\n",
    "\n",
    "axes[0].set_title(\"default grid\")\n",
    "axes[0].grid()\n",
    "\n",
    "axes[1].set_title(\"major/minor grid\")\n",
    "axes[1].grid(color=\"blue\", which=\"both\", linestyle=':', linewidth=0.5)\n",
    "\n",
    "axes[2].set_title(\"individual x/y major/minor grid\")\n",
    "axes[2].grid(color=\"grey\", which=\"major\", axis='x', linestyle='-', linewidth=0.5)\n",
    "axes[2].grid(color=\"grey\", which=\"minor\", axis='x', linestyle=':', linewidth=0.25)\n",
    "axes[2].grid(color=\"grey\", which=\"major\", axis='y', linestyle='-', linewidth=0.5)\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-axis-grid.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Ticker formatting"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 2, figsize=(8, 3))\n",
    "\n",
    "x = np.linspace(0, 1e5, 100)\n",
    "y = x ** 2\n",
    "\n",
    "axes[0].plot(x, y, 'b.')\n",
    "axes[0].set_title(\"default labels\", loc='right')\n",
    "\n",
    "axes[1].plot(x, y, 'b')\n",
    "axes[1].set_title(\"scientific notation labels\", loc='right')\n",
    "\n",
    "formatter = mpl.ticker.ScalarFormatter(useMathText=True)\n",
    "formatter.set_scientific(True) \n",
    "formatter.set_powerlimits((-1,1)) \n",
    "axes[1].xaxis.set_major_formatter(formatter)\n",
    "axes[1].yaxis.set_major_formatter(formatter)\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-axis-scientific.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Log plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(12, 3))\n",
    "\n",
    "x = np.linspace(0, 1e3, 100)\n",
    "y1, y2 = x**3, x**4\n",
    "\n",
    "axes[0].set_title('loglog')\n",
    "axes[0].loglog(x, y1, 'b', x, y2, 'r')\n",
    "\n",
    "axes[1].set_title('semilogy')\n",
    "axes[1].semilogy(x, y1, 'b', x, y2, 'r')\n",
    "\n",
    "axes[2].set_title('plot / set_xscale / set_yscale')\n",
    "axes[2].plot(x, y1, 'b', x, y2, 'r')\n",
    "axes[2].set_xscale('log')\n",
    "axes[2].set_yscale('log')\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-axis-log-plots.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Twin axes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax1 = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "r = np.linspace(0, 5, 100)\n",
    "a = 4 * np.pi * r ** 2  # area\n",
    "v = (4 * np.pi / 3) * r ** 3  # volume\n",
    "\n",
    "\n",
    "ax1.set_title(\"surface area and volume of a sphere\", fontsize=16)\n",
    "ax1.set_xlabel(\"radius [m]\", fontsize=16)\n",
    "\n",
    "ax1.plot(r, a, lw=2, color=\"blue\")\n",
    "ax1.set_ylabel(r\"surface area ($m^2$)\", fontsize=16, color=\"blue\")\n",
    "for label in ax1.get_yticklabels():\n",
    "    label.set_color(\"blue\")\n",
    "    \n",
    "ax2 = ax1.twinx()\n",
    "ax2.plot(r, v, lw=2, color=\"red\")\n",
    "ax2.set_ylabel(r\"volume ($m^3$)\", fontsize=16, color=\"red\")\n",
    "for label in ax2.get_yticklabels():\n",
    "    label.set_color(\"red\")\n",
    "    \n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-axis-twin-ax.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Spines"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = np.linspace(-10, 10, 500)\n",
    "y = np.sin(x) / x\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(8, 4))\n",
    "\n",
    "ax.plot(x, y, linewidth=2)\n",
    "\n",
    "# remove top and right spines\n",
    "ax.spines['right'].set_color('none')\n",
    "ax.spines['top'].set_color('none')\n",
    "\n",
    "# remove top and right spine ticks\n",
    "ax.xaxis.set_ticks_position('bottom')\n",
    "ax.yaxis.set_ticks_position('left')\n",
    "\n",
    "# move bottom and left spine to x = 0 and y = 0\n",
    "ax.spines['bottom'].set_position(('data', 0))\n",
    "ax.spines['left'].set_position(('data', 0))\n",
    "    \n",
    "ax.set_xticks([-10, -5, 5, 10])\n",
    "ax.set_yticks([0.5, 1])\n",
    "\n",
    "# give each label a solid background of white, to not overlap with the plot line\n",
    "for label in ax.get_xticklabels() + ax.get_yticklabels():\n",
    "    label.set_bbox({'facecolor': 'white',\n",
    "                    'edgecolor': 'white'})\n",
    "    \n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-axis-spines.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Advanced grid layout"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(8, 4))\n",
    "\n",
    "def f(x):\n",
    "    return 1/(1 + x**2) + 0.1/(1 + ((3 - x)/0.1)**2)\n",
    "\n",
    "def plot_and_format_axes(ax, x, f, fontsize):\n",
    "    ax.plot(x, f(x), linewidth=2)\n",
    "    ax.xaxis.set_major_locator(mpl.ticker.MaxNLocator(5))\n",
    "    ax.yaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "    ax.set_xlabel(r\"$x$\", fontsize=fontsize)\n",
    "    ax.set_ylabel(r\"$f(x)$\", fontsize=fontsize)\n",
    "    \n",
    "# main graph\n",
    "ax = fig.add_axes([0.1, 0.15, 0.8, 0.8], facecolor=\"#f5f5f5\")\n",
    "x = np.linspace(-4, 14, 1000)\n",
    "plot_and_format_axes(ax, x, f, 18)\n",
    "\n",
    "# inset\n",
    "x0, x1 = 2.5, 3.5\n",
    "ax.axvline(x0, ymax=0.3, color=\"grey\", linestyle=\":\")\n",
    "ax.axvline(x1, ymax=0.3, color=\"grey\", linestyle=\":\")\n",
    "\n",
    "ax = fig.add_axes([0.5, 0.5, 0.38, 0.42], facecolor='none')\n",
    "x = np.linspace(x0, x1, 1000)\n",
    "plot_and_format_axes(ax, x, f, 14)\n",
    "\n",
    "fig.savefig(\"ch4-advanced-axes-inset.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Subplots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "ncols, nrows = 3, 3\n",
    "\n",
    "fig, axes = plt.subplots(nrows, ncols)\n",
    "\n",
    "for m in range(nrows):\n",
    "    for n in range(ncols):\n",
    "        axes[m, n].set_xticks([])\n",
    "        axes[m, n].set_yticks([])\n",
    "        axes[m, n].text(0.5, 0.5, \"axes[%d, %d]\" % (m, n),\n",
    "                        horizontalalignment='center')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 2, figsize=(6, 6), sharex=True, sharey=True, squeeze=False)\n",
    "\n",
    "x1 = np.random.randn(100)\n",
    "x2 = np.random.randn(100)\n",
    "\n",
    "axes[0, 0].set_title(\"Uncorrelated\")\n",
    "axes[0, 0].scatter(x1, x2)\n",
    "\n",
    "axes[0, 1].set_title(\"Weakly positively correlated\")\n",
    "axes[0, 1].scatter(x1, x1 + x2)\n",
    "\n",
    "axes[1, 0].set_title(\"Weakly negatively correlated\")\n",
    "axes[1, 0].scatter(x1, -x1 + x2)\n",
    "\n",
    "axes[1, 1].set_title(\"Strongly correlated\")\n",
    "axes[1, 1].scatter(x1, x1 + 0.15 * x2)\n",
    "\n",
    "axes[1, 1].set_xlabel(\"x\")\n",
    "axes[1, 0].set_xlabel(\"x\")\n",
    "axes[0, 0].set_ylabel(\"y\")\n",
    "axes[1, 0].set_ylabel(\"y\")\n",
    "\n",
    "plt.subplots_adjust(left=0.1, right=0.95, bottom=0.1, top=0.95, wspace=0.1, hspace=0.2)\n",
    "\n",
    "fig.savefig(\"ch4-advanced-axes-subplots.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "\n",
    "def clear_ticklabels(ax):\n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_xticklabels([])\n",
    "\n",
    "ax0 = plt.subplot2grid((3, 3), (0, 0))\n",
    "ax1 = plt.subplot2grid((3, 3), (0, 1))\n",
    "ax2 = plt.subplot2grid((3, 3), (1, 0), colspan=2)\n",
    "ax3 = plt.subplot2grid((3, 3), (2, 0), colspan=3)\n",
    "ax4 = plt.subplot2grid((3, 3), (0, 2), rowspan=2)\n",
    "\n",
    "axes = [ax0, ax1, ax2, ax3, ax4]\n",
    "[ax.text(0.5, 0.5, \"ax%d\" % n, horizontalalignment='center') for n, ax in enumerate(axes)]\n",
    "[clear_ticklabels(ax) for ax in axes]\n",
    "\n",
    "fig.savefig(\"ch4-advanced-axes-subplot2grid.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### gridspec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from matplotlib.gridspec import GridSpec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(6, 4))\n",
    "\n",
    "gs = mpl.gridspec.GridSpec(4, 4)\n",
    "\n",
    "ax0 = fig.add_subplot(gs[0, 0])\n",
    "ax1 = fig.add_subplot(gs[1, 1])\n",
    "ax2 = fig.add_subplot(gs[2, 2])\n",
    "ax3 = fig.add_subplot(gs[3, 3])\n",
    "\n",
    "ax4 = fig.add_subplot(gs[0, 1:])\n",
    "ax5 = fig.add_subplot(gs[1:, 0])\n",
    "\n",
    "ax6 = fig.add_subplot(gs[1, 2:])\n",
    "ax7 = fig.add_subplot(gs[2:, 1])\n",
    "\n",
    "ax8 = fig.add_subplot(gs[2, 3])\n",
    "ax9 = fig.add_subplot(gs[3, 2])\n",
    "\n",
    "\n",
    "def clear_ticklabels(ax):\n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_xticklabels([])\n",
    "\n",
    "axes = [ax0, ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9]\n",
    "[ax.text(0.5, 0.5, \"ax%d\" % n, horizontalalignment='center') for n, ax in enumerate(axes)]\n",
    "[clear_ticklabels(ax) for ax in axes]\n",
    "\n",
    "fig.savefig(\"ch4-advanced-axes-gridspec-1.pdf\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(4, 4))\n",
    "\n",
    "gs = mpl.gridspec.GridSpec(2, 2,\n",
    "                           width_ratios=[4, 1],\n",
    "                           height_ratios=[1, 4],\n",
    "                           wspace=0.05, hspace=0.05\n",
    "                           )\n",
    "\n",
    "ax0 = fig.add_subplot(gs[1, 0])\n",
    "ax1 = fig.add_subplot(gs[0, 0])\n",
    "ax2 = fig.add_subplot(gs[1, 1])\n",
    "\n",
    "def clear_ticklabels(ax):\n",
    "    ax.set_yticklabels([])\n",
    "    ax.set_xticklabels([])\n",
    "\n",
    "axes = [ax0, ax1, ax2]\n",
    "[ax.text(0.5, 0.5, \"ax%d\" % n, horizontalalignment='center') for n, ax in enumerate(axes)]\n",
    "[clear_ticklabels(ax) for ax in axes]\n",
    "\n",
    "fig.savefig(\"ch4-advanced-axes-gridspec-2.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Colormap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = y = np.linspace(-2, 2, 150)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "\n",
    "R1 = np.sqrt((X+0.5)**2 + (Y+0.5)**2)\n",
    "R2 = np.sqrt((X+0.5)**2 + (Y-0.5)**2)\n",
    "R3 = np.sqrt((X-0.5)**2 + (Y+0.5)**2)\n",
    "R4 = np.sqrt((X-0.5)**2 + (Y-0.5)**2)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Z = np.sin(10 * R1) / (10 * R1) + np.sin(20 * R4) / (20 * R4)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 5))\n",
    "\n",
    "p = ax.pcolor(X, Y, Z, cmap='seismic', vmin=-abs(Z).max(), vmax=abs(Z).max())\n",
    "ax.axis('tight')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "cb = fig.colorbar(p, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "Z = 1/R1 - 1/R2 - 1/R3 + 1/R4\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 5))\n",
    "\n",
    "im = ax.imshow(Z, vmin=-1, vmax=1, cmap=mpl.cm.bwr,\n",
    "               extent=[x.min(), x.max(), y.min(), y.max()])\n",
    "im.set_interpolation('bilinear')\n",
    "\n",
    "ax.axis('tight')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')\n",
    "cb = fig.colorbar(p, ax=ax)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = y = np.linspace(-2, 2, 150)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "\n",
    "R1 = np.sqrt((X+0.5)**2 + (Y+0.5)**2)\n",
    "R2 = np.sqrt((X+0.5)**2 + (Y-0.5)**2)\n",
    "R3 = np.sqrt((X-0.5)**2 + (Y+0.5)**2)\n",
    "R4 = np.sqrt((X-0.5)**2 + (Y-0.5)**2)\n",
    "\n",
    "fig, axes = plt.subplots(1, 4, figsize=(14, 3))\n",
    "\n",
    "Z = np.sin(10 * R1) / (10 * R1) + np.sin(20 * R4) / (20 * R4)\n",
    "\n",
    "\n",
    "p = axes[0].pcolor(X, Y, Z, cmap='seismic', vmin=-abs(Z).max(), vmax=abs(Z).max())\n",
    "axes[0].axis('tight')\n",
    "axes[0].set_xlabel('x')\n",
    "axes[0].set_ylabel('y')\n",
    "axes[0].set_title(\"pcolor\")\n",
    "axes[0].xaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "axes[0].yaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "\n",
    "\n",
    "cb = fig.colorbar(p, ax=axes[0])\n",
    "cb.set_label(\"z\")\n",
    "cb.set_ticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "\n",
    "Z = 1/R1 - 1/R2 - 1/R3 + 1/R4\n",
    "\n",
    "im = axes[1].imshow(Z, vmin=-1, vmax=1, cmap=mpl.cm.bwr,\n",
    "               extent=[x.min(), x.max(), y.min(), y.max()])\n",
    "im.set_interpolation('bilinear')\n",
    "\n",
    "axes[1].axis('tight')\n",
    "axes[1].set_xlabel('x')\n",
    "axes[1].set_ylabel('y')\n",
    "axes[1].set_title(\"imshow\")\n",
    "cb = fig.colorbar(im, ax=axes[1])\n",
    "\n",
    "axes[1].xaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "axes[1].yaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "#cb.ax.set_axes_locator(mpl.ticker.MaxNLocator(4))\n",
    "cb.set_label(\"z\")\n",
    "cb.set_ticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "x = y = np.linspace(0, 1, 75)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "Z = - 2 * np.cos(2 * np.pi * X) * np.cos(2 * np.pi * Y) - 0.7 * np.cos(np.pi - 4*np.pi*X)\n",
    "\n",
    "c = axes[2].contour(X, Y, Z, 15, cmap=mpl.cm.RdBu, vmin=-1, vmax=1)\n",
    "\n",
    "axes[2].axis('tight')\n",
    "axes[2].set_xlabel('x')\n",
    "axes[2].set_ylabel('y')\n",
    "axes[2].set_title(\"contour\")\n",
    "\n",
    "axes[2].xaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "axes[2].yaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "\n",
    "\n",
    "c = axes[3].contourf(X, Y, Z, 15, cmap=mpl.cm.RdBu, vmin=-1, vmax=1)\n",
    "\n",
    "axes[3].axis('tight')\n",
    "axes[3].set_xlabel('x')\n",
    "axes[3].set_ylabel('y')\n",
    "axes[3].set_title(\"contourf\")\n",
    "\n",
    "axes[3].xaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "axes[3].yaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig('ch4-colormaps.pdf')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots(figsize=(6, 5))\n",
    "\n",
    "x = y = np.linspace(0, 1, 75)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "\n",
    "Z = - 2 * np.cos(2 * np.pi * X) * np.cos(2 * np.pi * Y) - 0.7 * np.cos(np.pi - 4*np.pi*X)\n",
    "\n",
    "c = ax.contour(X, Y, Z, 15, cmap=mpl.cm.RdBu, vmin=-1, vmax=1)\n",
    "\n",
    "ax.axis('tight')\n",
    "ax.set_xlabel('x')\n",
    "ax.set_ylabel('y')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = y = np.linspace(-10, 10, 150)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "Z = np.cos(X) * np.cos(Y) * np.exp(-(X/5)**2-(Y/5)**2)\n",
    "\n",
    "fig, ax = plt.subplots(figsize=(6, 5))\n",
    "\n",
    "p = ax.pcolor(X, Y, Z, vmin=-abs(Z).max(), vmax=abs(Z).max(), cmap=mpl.cm.bwr)\n",
    "\n",
    "ax.axis('tight')\n",
    "ax.set_xlabel(r\"$x$\", fontsize=18)\n",
    "ax.set_ylabel(r\"$y$\", fontsize=18)\n",
    "ax.xaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "ax.yaxis.set_major_locator(mpl.ticker.MaxNLocator(4))\n",
    "\n",
    "cb = fig.colorbar(p, ax=ax)\n",
    "cb.set_label(r\"$z$\", fontsize=18)\n",
    "cb.set_ticks([-1, -.5, 0, .5, 1])\n",
    "\n",
    "fig.savefig(\"ch4-colormap-pcolor.pdf\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3D plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from mpl_toolkits.mplot3d.axes3d import Axes3D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "x = y = np.linspace(-3, 3, 74)\n",
    "X, Y = np.meshgrid(x, y)\n",
    "\n",
    "R = np.sqrt(X**2 + Y**2)\n",
    "Z = np.sin(4 * R) / R"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(14, 4), subplot_kw={'projection': '3d'})\n",
    "\n",
    "def title_and_labels(ax, title):\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"$x$\", fontsize=16)\n",
    "    ax.set_ylabel(\"$y$\", fontsize=16)\n",
    "    ax.set_zlabel(\"$z$\", fontsize=16)\n",
    "\n",
    "norm = mpl.colors.Normalize(-abs(Z).max(), abs(Z).max())\n",
    "\n",
    "p = axes[0].plot_surface(X, Y, Z, rstride=1, cstride=1, linewidth=0, antialiased=False, norm=norm, cmap=mpl.cm.Blues)\n",
    "cb = fig.colorbar(p, ax=axes[0], shrink=0.6)\n",
    "title_and_labels(axes[0], \"plot_surface\")\n",
    "\n",
    "p = axes[1].plot_wireframe(X, Y, Z, rstride=2, cstride=2, color=\"darkgrey\")\n",
    "title_and_labels(axes[1], \"plot_wireframe\")\n",
    "\n",
    "cset = axes[2].contour(X, Y, Z, zdir='z', offset=0, norm=norm, cmap=mpl.cm.Blues)\n",
    "cset = axes[2].contour(X, Y, Z, zdir='y', offset=3, norm=norm, cmap=mpl.cm.Blues)\n",
    "title_and_labels(axes[2], \"contour\")\n",
    "\n",
    "fig.tight_layout()\n",
    "fig.savefig(\"ch4-3d-plots.png\", dpi=200)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(1, 3, figsize=(14, 4), subplot_kw={'projection': '3d'})\n",
    "\n",
    "def title_and_labels(ax, title):\n",
    "    ax.set_title(title)\n",
    "    ax.set_xlabel(\"$x$\", fontsize=16)\n",
    "    ax.set_ylabel(\"$y$\", fontsize=16)\n",
    "    ax.set_zlabel(\"$z$\", fontsize=16)\n",
    "\n",
    "norm = mpl.colors.Normalize(-abs(Z).max(), abs(Z).max())\n",
    "\n",
    "r = np.linspace(0, 10, 100)\n",
    "p = axes[0].plot(np.cos(r), np.sin(r), 6 - r)\n",
    "#cb = fig.colorbar(p, ax=axes[0], shrink=0.6)\n",
    "title_and_labels(axes[0], \"plot\")\n",
    "\n",
    "p = axes[1].scatter(np.cos(r), np.sin(r), 6 - r)\n",
    "title_and_labels(axes[1], \"scatter\")\n",
    "\n",
    "r = np.linspace(0, 6, 20)\n",
    "p = axes[2].bar3d(np.cos(r), np.sin(r), 0* np.ones_like(r), 0.05* np.ones_like(r), 0.05 * np.ones_like(r), 6 - r)\n",
    "title_and_labels(axes[2], \"contour\")\n",
    "axes[2].text(0, 0, 0, \"label\")\n",
    "fig.tight_layout()\n",
    "#fig.savefig(\"ch4-3d-plots.png\", dpi=200)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "%reload_ext version_information"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%version_information numpy, matplotlib"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.6 (NPM ed2)",
   "language": "python",
   "name": "py3.6-npm-e2"
  },
  "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.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}