{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4502c531",
   "metadata": {},
   "source": [
    "\n",
    "<a id='np'></a>\n",
    "<div id=\"qe-notebook-header\" align=\"right\" style=\"text-align:right;\">\n",
    "        <a href=\"https://quantecon.org/\" title=\"quantecon.org\">\n",
    "                <img style=\"width:250px;display:inline;\" width=\"250px\" src=\"https://assets.quantecon.org/img/qe-menubar-logo.svg\" alt=\"QuantEcon\">\n",
    "        </a>\n",
    "</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65cf7149",
   "metadata": {},
   "source": [
    "# NumPy\n",
    "\n",
    "\n",
    "<a id='index-1'></a>\n",
    "> “Let’s be clear: the work of science has nothing whatever to do with consensus.  Consensus is the business of politics. Science, on the contrary, requires only one investigator who happens to be right, which means that he or she has results that are verifiable by reference to the real world. In science consensus is irrelevant. What is relevant is reproducible results.” – Michael Crichton\n",
    "\n",
    "\n",
    "In addition to what’s in Anaconda, this lecture will need the following libraries:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87acde12",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63855327",
   "metadata": {},
   "source": [
    "## Overview\n",
    "\n",
    "[NumPy](https://en.wikipedia.org/wiki/NumPy) is a first-rate library for numerical programming\n",
    "\n",
    "- Widely used in academia, finance and industry.  \n",
    "- Mature, fast, stable and under continuous development.  \n",
    "\n",
    "\n",
    "We have already seen some code involving NumPy in the preceding lectures.\n",
    "\n",
    "In this lecture, we will start a more systematic discussion of\n",
    "\n",
    "1. NumPy arrays and  \n",
    "1. the fundamental array processing operations provided by NumPy.  \n",
    "\n",
    "\n",
    "(For an alternative reference, see [the official NumPy documentation](https://docs.scipy.org/doc/numpy/reference/).)\n",
    "\n",
    "We will use the following imports."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1fdfbd2",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import random\n",
    "import quantecon as qe\n",
    "import matplotlib.pyplot as plt\n",
    "from mpl_toolkits.mplot3d.axes3d import Axes3D\n",
    "from matplotlib import cm"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afa138b1",
   "metadata": {},
   "source": [
    "\n",
    "<a id='numpy-array'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62e0a2a2",
   "metadata": {},
   "source": [
    "## NumPy Arrays\n",
    "\n",
    "\n",
    "<a id='index-2'></a>\n",
    "The essential problem that NumPy solves is fast array processing.\n",
    "\n",
    "The most important structure that NumPy defines is an array data type, formally\n",
    "called a [numpy.ndarray](https://docs.scipy.org/doc/numpy/reference/arrays.ndarray.html).\n",
    "\n",
    "NumPy arrays power a very large proportion of the scientific Python ecosystem.\n",
    "\n",
    "To create a NumPy array containing only zeros we use  [np.zeros](https://docs.scipy.org/doc/numpy/reference/generated/numpy.zeros.html#numpy.zeros)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1978d6e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.zeros(3)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a4144991",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "type(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "03a16ffc",
   "metadata": {},
   "source": [
    "NumPy arrays are somewhat like native Python lists, except that\n",
    "\n",
    "- Data *must be homogeneous* (all elements of the same type).  \n",
    "- These types must be one of the [data types](https://docs.scipy.org/doc/numpy/reference/arrays.dtypes.html) (`dtypes`) provided by NumPy.  \n",
    "\n",
    "\n",
    "The most important of these dtypes are:\n",
    "\n",
    "- float64: 64 bit floating-point number  \n",
    "- int64: 64 bit integer  \n",
    "- bool:  8 bit True or False  \n",
    "\n",
    "\n",
    "There are also dtypes to represent complex numbers, unsigned integers, etc.\n",
    "\n",
    "On modern machines, the default dtype for arrays is `float64`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a5fcbce",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.zeros(3)\n",
    "type(a[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16b6a990",
   "metadata": {},
   "source": [
    "If we want to use integers we can specify as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a7823de",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.zeros(3, dtype=int)\n",
    "type(a[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "964660bb",
   "metadata": {},
   "source": [
    "\n",
    "<a id='numpy-shape-dim'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9301148d",
   "metadata": {},
   "source": [
    "### Shape and Dimension\n",
    "\n",
    "\n",
    "<a id='index-3'></a>\n",
    "Consider the following assignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6183d729",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.zeros(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5baa70b6",
   "metadata": {},
   "source": [
    "Here `z` is a *flat* array with no dimension — neither row nor column vector.\n",
    "\n",
    "The dimension is recorded in the `shape` attribute, which is a tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f46753fb",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48ab606c",
   "metadata": {},
   "source": [
    "Here the shape tuple has only one element, which is the length of the array (tuples with one element end with a comma).\n",
    "\n",
    "To give it dimension, we can change the `shape` attribute"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e2f43d3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z.shape = (10, 1)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04d84d26",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.zeros(4)\n",
    "z.shape = (2, 2)\n",
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dbfb66ee",
   "metadata": {},
   "source": [
    "In the last case, to make the 2 by 2 array, we could also pass a tuple to the `zeros()` function, as\n",
    "in `z = np.zeros((2, 2))`.\n",
    "\n",
    "\n",
    "<a id='creating-arrays'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f45a24df",
   "metadata": {},
   "source": [
    "### Creating Arrays\n",
    "\n",
    "\n",
    "<a id='index-4'></a>\n",
    "As we’ve seen, the `np.zeros` function creates an array of zeros.\n",
    "\n",
    "You can probably guess what `np.ones` creates.\n",
    "\n",
    "Related is `np.empty`, which creates arrays in memory that can later be populated with data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60c58574",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.empty(3)\n",
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "456112db",
   "metadata": {},
   "source": [
    "The numbers you see here are garbage values.\n",
    "\n",
    "(Python allocates 3 contiguous 64 bit pieces of memory, and the existing contents of those memory slots are interpreted as `float64` values)\n",
    "\n",
    "To set up a grid of evenly spaced numbers use `np.linspace`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d8a14ca",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.linspace(2, 4, 5)  # From 2 to 4, with 5 elements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b37e58cf",
   "metadata": {},
   "source": [
    "To create an identity matrix use either `np.identity` or `np.eye`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "54249464",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.identity(2)\n",
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aacc327a",
   "metadata": {},
   "source": [
    "In addition, NumPy arrays can be created from Python lists, tuples, etc. using `np.array`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd5d1556",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.array([10, 20])                 # ndarray from Python list\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f792eb0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "type(z)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1d48200",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.array((10, 20), dtype=float)    # Here 'float' is equivalent to 'np.float64'\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2e20fce9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.array([[1, 2], [3, 4]])         # 2D array from a list of lists\n",
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5145c8a",
   "metadata": {},
   "source": [
    "See also `np.asarray`, which performs a similar function, but does not make\n",
    "a distinct copy of data already in a NumPy array."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61d9d6fc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "na = np.linspace(10, 20, 2)\n",
    "na is np.asarray(na)   # Does not copy NumPy arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf61750f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "na is np.array(na)     # Does make a new copy --- perhaps unnecessarily"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "28dcb02a",
   "metadata": {},
   "source": [
    "To read in the array data from a text file containing numeric data use `np.loadtxt`\n",
    "or `np.genfromtxt`—see [the documentation](https://docs.scipy.org/doc/numpy/reference/routines.io.html) for details."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00f09862",
   "metadata": {},
   "source": [
    "### Array Indexing\n",
    "\n",
    "\n",
    "<a id='index-5'></a>\n",
    "For a flat array, indexing is the same as Python sequences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72313d55",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.linspace(1, 2, 5)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a565435",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1d6ad00d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[0:2]  # Two elements, starting at element 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c46644c7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[-1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a51381ee",
   "metadata": {},
   "source": [
    "For 2D arrays the index syntax is as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b6c2393",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.array([[1, 2], [3, 4]])\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e131e8a8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[0, 0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "03f6fc7c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[0, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b949d15",
   "metadata": {},
   "source": [
    "And so on.\n",
    "\n",
    "Note that indices are still zero-based, to maintain compatibility with Python sequences.\n",
    "\n",
    "Columns and rows can be extracted as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6a231b7e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[0, :]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65d56f9e",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[:, 1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7692cae4",
   "metadata": {},
   "source": [
    "NumPy arrays of integers can also be used to extract elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e2808fa",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.linspace(2, 4, 5)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1d2b6ef",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "indices = np.array((0, 2, 3))\n",
    "z[indices]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c03d79e",
   "metadata": {},
   "source": [
    "Finally, an array of `dtype bool` can be used to extract elements"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "452302e1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873c431a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "d = np.array([0, 1, 1, 0, 0], dtype=bool)\n",
    "d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c7bb84",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[d]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a13330cd",
   "metadata": {},
   "source": [
    "We’ll see why this is useful below.\n",
    "\n",
    "An aside: all elements of an array can be set equal to one number using slice notation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cdb1b0c5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.empty(3)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5fb054ed",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[:] = 42\n",
    "z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10b36de9",
   "metadata": {},
   "source": [
    "### Array Methods\n",
    "\n",
    "\n",
    "<a id='index-6'></a>\n",
    "Arrays have useful methods, all of which are carefully optimized"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "295e720c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array((4, 3, 2, 1))\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e78a8c44",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.sort()              # Sorts a in place\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e574280a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.sum()               # Sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be655701",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.mean()              # Mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a04295d6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.max()               # Max"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ee16ff6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.argmax()            # Returns the index of the maximal element"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "206b7fbe",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.cumsum()            # Cumulative sum of the elements of a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bc47f39",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.cumprod()           # Cumulative product of the elements of a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f593c33",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.var()               # Variance"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4476075f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.std()               # Standard deviation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc590e8c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a.shape = (2, 2)\n",
    "a.T                   # Equivalent to a.transpose()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db8f5b8c",
   "metadata": {},
   "source": [
    "Another method worth knowing is `searchsorted()`.\n",
    "\n",
    "If `z` is a nondecreasing array, then `z.searchsorted(a)` returns the index of the first element of `z` that is `>= a`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c0a23c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.linspace(2, 4, 5)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b539b6af",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z.searchsorted(2.2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46d8058f",
   "metadata": {},
   "source": [
    "Many of the methods discussed above have equivalent functions in the NumPy namespace"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f41e2dc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array((4, 3, 2, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fbb0c791",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.sum(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba602c5a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.mean(a)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d0462c2",
   "metadata": {},
   "source": [
    "## Arithmetic Operations\n",
    "\n",
    "\n",
    "<a id='index-7'></a>\n",
    "The operators `+`, `-`, `*`, `/` and `**` all act *elementwise* on arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aa0a6381",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array([1, 2, 3, 4])\n",
    "b = np.array([5, 6, 7, 8])\n",
    "a + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2335340a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a * b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f95ecd5",
   "metadata": {},
   "source": [
    "We can add a scalar to each element as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ebaef052",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a + 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74126ad8",
   "metadata": {},
   "source": [
    "Scalar multiplication is similar"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ec30a6f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a * 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf8afc7b",
   "metadata": {},
   "source": [
    "The two-dimensional arrays follow the same general rules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2c263fe",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.ones((2, 2))\n",
    "B = np.ones((2, 2))\n",
    "A + B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24039639",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A + 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82304a42",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A * B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7c5302e",
   "metadata": {},
   "source": [
    "\n",
    "<a id='numpy-matrix-multiplication'></a>\n",
    "In particular, `A * B` is *not* the matrix product, it is an element-wise product."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f739f742",
   "metadata": {},
   "source": [
    "## Matrix Multiplication\n",
    "\n",
    "\n",
    "<a id='index-8'></a>\n",
    "\n",
    "<a id='index-9'></a>\n",
    "With Anaconda’s scientific Python package based around Python 3.5 and above,\n",
    "one can use the `@` symbol for matrix multiplication, as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b0880c6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.ones((2, 2))\n",
    "B = np.ones((2, 2))\n",
    "A @ B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d192ecd5",
   "metadata": {},
   "source": [
    "(For older versions of Python and NumPy you need to use the [np.dot](https://docs.scipy.org/doc/numpy/reference/generated/numpy.dot.html) function)\n",
    "\n",
    "We can also use `@` to take the inner product of two flat arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2e59f71",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.array((1, 2))\n",
    "B = np.array((10, 20))\n",
    "A @ B"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dab1ab63",
   "metadata": {},
   "source": [
    "In fact, we can use `@` when one element is a Python list or tuple"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc483ea4",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.array(((1, 2), (3, 4)))\n",
    "A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "881d4677",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A @ (0, 1)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b85dcbe",
   "metadata": {},
   "source": [
    "Since we are post-multiplying, the tuple is treated as a column vector.\n",
    "\n",
    "\n",
    "<a id='broadcasting'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c7b947f",
   "metadata": {},
   "source": [
    "## Broadcasting\n",
    "\n",
    "\n",
    "<a id='index-10'></a>\n",
    "(This section extends an excellent discussion of broadcasting provided by [Jake VanderPlas](https://jakevdp.github.io/PythonDataScienceHandbook/02.05-computation-on-arrays-broadcasting.html).)\n",
    "\n",
    ">**Note**\n",
    ">\n",
    ">Broadcasting is a very important aspect of NumPy. At the same time, advanced broadcasting is relatively complex and some of the details below can be skimmed on first pass.\n",
    "\n",
    "In element-wise operations, arrays may not have the same shape.\n",
    "\n",
    "When this happens, NumPy will automatically expand arrays to the same shape whenever possible.\n",
    "\n",
    "This useful (but sometimes confusing) feature in NumPy is called **broadcasting**.\n",
    "\n",
    "The value of broadcasting is that\n",
    "\n",
    "- `for` loops can be avoided, which helps numerical code run fast and  \n",
    "- broadcasting can allow us to implement operations on arrays without actually creating some dimensions of these arrays in memory, which can be important when arrays are large.  \n",
    "\n",
    "\n",
    "For example, suppose `a` is a $ 3 \\times 3 $ array (`a -> (3, 3)`), while `b` is a flat array with three elements (`b -> (3,)`).\n",
    "\n",
    "When adding them together, NumPy will automatically expand `b -> (3,)` to `b -> (3, 3)`.\n",
    "\n",
    "The element-wise addition will result in a $ 3 \\times 3 $ array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9ef03d7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array(\n",
    "        [[1, 2, 3], \n",
    "         [4, 5, 6], \n",
    "         [7, 8, 9]])\n",
    "b = np.array([3, 6, 9])\n",
    "\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0758ac84",
   "metadata": {},
   "source": [
    "Here is a visual representation of this broadcasting operation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "976cbfee",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Adapted and modified based on the code in the book written by Jake VanderPlas (see https://jakevdp.github.io/PythonDataScienceHandbook/06.00-figure-code.html#Broadcasting)\n",
    "# Originally from astroML: see https://www.astroml.org/book_figures/appendix/fig_broadcast_visual.html\n",
    "\n",
    "\n",
    "def draw_cube(ax, xy, size, depth=0.4,\n",
    "              edges=None, label=None, label_kwargs=None, **kwargs):\n",
    "    \"\"\"draw and label a cube.  edges is a list of numbers between\n",
    "    1 and 12, specifying which of the 12 cube edges to draw\"\"\"\n",
    "    if edges is None:\n",
    "        edges = range(1, 13)\n",
    "\n",
    "    x, y = xy\n",
    "\n",
    "    if 1 in edges:\n",
    "        ax.plot([x, x + size],\n",
    "                [y + size, y + size], **kwargs)\n",
    "    if 2 in edges:\n",
    "        ax.plot([x + size, x + size],\n",
    "                [y, y + size], **kwargs)\n",
    "    if 3 in edges:\n",
    "        ax.plot([x, x + size],\n",
    "                [y, y], **kwargs)\n",
    "    if 4 in edges:\n",
    "        ax.plot([x, x],\n",
    "                [y, y + size], **kwargs)\n",
    "\n",
    "    if 5 in edges:\n",
    "        ax.plot([x, x + depth],\n",
    "                [y + size, y + depth + size], **kwargs)\n",
    "    if 6 in edges:\n",
    "        ax.plot([x + size, x + size + depth],\n",
    "                [y + size, y + depth + size], **kwargs)\n",
    "    if 7 in edges:\n",
    "        ax.plot([x + size, x + size + depth],\n",
    "                [y, y + depth], **kwargs)\n",
    "    if 8 in edges:\n",
    "        ax.plot([x, x + depth],\n",
    "                [y, y + depth], **kwargs)\n",
    "\n",
    "    if 9 in edges:\n",
    "        ax.plot([x + depth, x + depth + size],\n",
    "                [y + depth + size, y + depth + size], **kwargs)\n",
    "    if 10 in edges:\n",
    "        ax.plot([x + depth + size, x + depth + size],\n",
    "                [y + depth, y + depth + size], **kwargs)\n",
    "    if 11 in edges:\n",
    "        ax.plot([x + depth, x + depth + size],\n",
    "                [y + depth, y + depth], **kwargs)\n",
    "    if 12 in edges:\n",
    "        ax.plot([x + depth, x + depth],\n",
    "                [y + depth, y + depth + size], **kwargs)\n",
    "\n",
    "    if label:\n",
    "        if label_kwargs is None:\n",
    "            label_kwargs = {}\n",
    "        ax.text(x + 0.5 * size, y + 0.5 * size, label,\n",
    "                ha='center', va='center', **label_kwargs)\n",
    "\n",
    "solid = dict(c='black', ls='-', lw=1,\n",
    "             label_kwargs=dict(color='k'))\n",
    "dotted = dict(c='black', ls='-', lw=0.5, alpha=0.5,\n",
    "              label_kwargs=dict(color='gray'))\n",
    "depth = 0.3\n",
    "\n",
    "# Draw a figure and axis with no boundary\n",
    "fig = plt.figure(figsize=(5, 1), facecolor='w')\n",
    "ax = plt.axes([0, 0, 1, 1], xticks=[], yticks=[], frameon=False)\n",
    "\n",
    "# first block\n",
    "draw_cube(ax, (1, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '1', **solid)\n",
    "draw_cube(ax, (2, 7.5), 1, depth, [1, 2, 3, 6, 9], '2', **solid)\n",
    "draw_cube(ax, (3, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '3', **solid)\n",
    "\n",
    "draw_cube(ax, (1, 6.5), 1, depth, [2, 3, 4], '4', **solid)\n",
    "draw_cube(ax, (2, 6.5), 1, depth, [2, 3], '5', **solid)\n",
    "draw_cube(ax, (3, 6.5), 1, depth, [2, 3, 7, 10], '6', **solid)\n",
    "\n",
    "draw_cube(ax, (1, 5.5), 1, depth, [2, 3, 4], '7', **solid)\n",
    "draw_cube(ax, (2, 5.5), 1, depth, [2, 3], '8', **solid)\n",
    "draw_cube(ax, (3, 5.5), 1, depth, [2, 3, 7, 10], '9', **solid)\n",
    "\n",
    "# second block\n",
    "draw_cube(ax, (6, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '3', **solid)\n",
    "draw_cube(ax, (7, 7.5), 1, depth, [1, 2, 3, 6, 9], '6', **solid)\n",
    "draw_cube(ax, (8, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '9', **solid)\n",
    "\n",
    "draw_cube(ax, (6, 6.5), 1, depth, range(2, 13), '3', **dotted)\n",
    "draw_cube(ax, (7, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '6', **dotted)\n",
    "draw_cube(ax, (8, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '9', **dotted)\n",
    "\n",
    "draw_cube(ax, (6, 5.5), 1, depth, [2, 3, 4, 7, 8, 10, 11, 12], '3', **dotted)\n",
    "draw_cube(ax, (7, 5.5), 1, depth, [2, 3, 7, 10, 11], '6', **dotted)\n",
    "draw_cube(ax, (8, 5.5), 1, depth, [2, 3, 7, 10, 11], '9', **dotted)\n",
    "\n",
    "# third block\n",
    "draw_cube(ax, (12, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '4', **solid)\n",
    "draw_cube(ax, (13, 7.5), 1, depth, [1, 2, 3, 6, 9], '8', **solid)\n",
    "draw_cube(ax, (14, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '12', **solid)\n",
    "\n",
    "draw_cube(ax, (12, 6.5), 1, depth, [2, 3, 4], '7', **solid)\n",
    "draw_cube(ax, (13, 6.5), 1, depth, [2, 3], '11', **solid)\n",
    "draw_cube(ax, (14, 6.5), 1, depth, [2, 3, 7, 10], '15', **solid)\n",
    "\n",
    "draw_cube(ax, (12, 5.5), 1, depth, [2, 3, 4], '10', **solid)\n",
    "draw_cube(ax, (13, 5.5), 1, depth, [2, 3], '14', **solid)\n",
    "draw_cube(ax, (14, 5.5), 1, depth, [2, 3, 7, 10], '18', **solid)\n",
    "\n",
    "ax.text(5, 7.0, '+', size=12, ha='center', va='center')\n",
    "ax.text(10.5, 7.0, '=', size=12, ha='center', va='center');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d087bdfa",
   "metadata": {},
   "source": [
    "How about `b -> (3, 1)`?\n",
    "\n",
    "In this case, NumPy will automatically expand `b -> (3, 1)` to `b -> (3, 3)`.\n",
    "\n",
    "Element-wise addition will then result in a $ 3 \\times 3 $ matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8a867d7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "b.shape = (3, 1)\n",
    "\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d206ccea",
   "metadata": {},
   "source": [
    "Here is a visual representation of this broadcasting operation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70eaa638",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig = plt.figure(figsize=(5, 1), facecolor='w')\n",
    "ax = plt.axes([0, 0, 1, 1], xticks=[], yticks=[], frameon=False)\n",
    "\n",
    "# first block\n",
    "draw_cube(ax, (1, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '1', **solid)\n",
    "draw_cube(ax, (2, 7.5), 1, depth, [1, 2, 3, 6, 9], '2', **solid)\n",
    "draw_cube(ax, (3, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '3', **solid)\n",
    "\n",
    "draw_cube(ax, (1, 6.5), 1, depth, [2, 3, 4], '4', **solid)\n",
    "draw_cube(ax, (2, 6.5), 1, depth, [2, 3], '5', **solid)\n",
    "draw_cube(ax, (3, 6.5), 1, depth, [2, 3, 7, 10], '6', **solid)\n",
    "\n",
    "draw_cube(ax, (1, 5.5), 1, depth, [2, 3, 4], '7', **solid)\n",
    "draw_cube(ax, (2, 5.5), 1, depth, [2, 3], '8', **solid)\n",
    "draw_cube(ax, (3, 5.5), 1, depth, [2, 3, 7, 10], '9', **solid)\n",
    "\n",
    "# second block\n",
    "draw_cube(ax, (6, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 7, 9, 10], '3', **solid)\n",
    "draw_cube(ax, (7, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '3', **dotted)\n",
    "draw_cube(ax, (8, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '3', **dotted)\n",
    "\n",
    "draw_cube(ax, (6, 6.5), 1, depth, [2, 3, 4, 7, 10], '6', **solid)\n",
    "draw_cube(ax, (7, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '6', **dotted)\n",
    "draw_cube(ax, (8, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '6', **dotted)\n",
    "\n",
    "draw_cube(ax, (6, 5.5), 1, depth, [2, 3, 4, 7, 10], '9', **solid)\n",
    "draw_cube(ax, (7, 5.5), 1, depth, [2, 3, 7, 10, 11], '9', **dotted)\n",
    "draw_cube(ax, (8, 5.5), 1, depth, [2, 3, 7, 10, 11], '9', **dotted)\n",
    "\n",
    "# third block\n",
    "draw_cube(ax, (12, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '4', **solid)\n",
    "draw_cube(ax, (13, 7.5), 1, depth, [1, 2, 3, 6, 9], '5', **solid)\n",
    "draw_cube(ax, (14, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '6', **solid)\n",
    "\n",
    "draw_cube(ax, (12, 6.5), 1, depth, [2, 3, 4], '10', **solid)\n",
    "draw_cube(ax, (13, 6.5), 1, depth, [2, 3], '11', **solid)\n",
    "draw_cube(ax, (14, 6.5), 1, depth, [2, 3, 7, 10], '12', **solid)\n",
    "\n",
    "draw_cube(ax, (12, 5.5), 1, depth, [2, 3, 4], '16', **solid)\n",
    "draw_cube(ax, (13, 5.5), 1, depth, [2, 3], '17', **solid)\n",
    "draw_cube(ax, (14, 5.5), 1, depth, [2, 3, 7, 10], '18', **solid)\n",
    "\n",
    "ax.text(5, 7.0, '+', size=12, ha='center', va='center')\n",
    "ax.text(10.5, 7.0, '=', size=12, ha='center', va='center');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "75b8b225",
   "metadata": {},
   "source": [
    "The previous broadcasting operation is equivalent to the following `for` loop"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aea7c86c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "row, column = a.shape\n",
    "result = np.empty((3, 3))\n",
    "for i in range(row):\n",
    "    for j in range(column):\n",
    "        result[i, j] = a[i, j] + b[i,0]\n",
    "\n",
    "result"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48da012d",
   "metadata": {},
   "source": [
    "In some cases, both operands will be expanded.\n",
    "\n",
    "When we have `a -> (3,)` and `b -> (3, 1)`, `a` will be expanded to `a -> (3, 3)`, and `b` will be expanded to `b -> (3, 3)`.\n",
    "\n",
    "In this case, element-wise addition will result in a $ 3 \\times 3 $ matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0599dca",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array([3, 6, 9])\n",
    "b = np.array([2, 3, 4])\n",
    "b.shape = (3, 1)\n",
    "\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9ededd5",
   "metadata": {},
   "source": [
    "Here is a visual representation of this broadcasting operation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f998bb8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Draw a figure and axis with no boundary\n",
    "fig = plt.figure(figsize=(5, 1), facecolor='w')\n",
    "ax = plt.axes([0, 0, 1, 1], xticks=[], yticks=[], frameon=False)\n",
    "\n",
    "# first block\n",
    "draw_cube(ax, (1, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '3', **solid)\n",
    "draw_cube(ax, (2, 7.5), 1, depth, [1, 2, 3, 6, 9], '6', **solid)\n",
    "draw_cube(ax, (3, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '9', **solid)\n",
    "\n",
    "draw_cube(ax, (1, 6.5), 1, depth, range(2, 13), '3', **dotted)\n",
    "draw_cube(ax, (2, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '6', **dotted)\n",
    "draw_cube(ax, (3, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '9', **dotted)\n",
    "\n",
    "draw_cube(ax, (1, 5.5), 1, depth, [2, 3, 4, 7, 8, 10, 11, 12], '3', **dotted)\n",
    "draw_cube(ax, (2, 5.5), 1, depth, [2, 3, 7, 10, 11], '6', **dotted)\n",
    "draw_cube(ax, (3, 5.5), 1, depth, [2, 3, 7, 10, 11], '9', **dotted)\n",
    "\n",
    "# second block\n",
    "draw_cube(ax, (6, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 7, 9, 10], '2', **solid)\n",
    "draw_cube(ax, (7, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '2', **dotted)\n",
    "draw_cube(ax, (8, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '2', **dotted)\n",
    "\n",
    "draw_cube(ax, (6, 6.5), 1, depth, [2, 3, 4, 7, 10], '3', **solid)\n",
    "draw_cube(ax, (7, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '3', **dotted)\n",
    "draw_cube(ax, (8, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '3', **dotted)\n",
    "\n",
    "draw_cube(ax, (6, 5.5), 1, depth, [2, 3, 4, 7, 10], '4', **solid)\n",
    "draw_cube(ax, (7, 5.5), 1, depth, [2, 3, 7, 10, 11], '4', **dotted)\n",
    "draw_cube(ax, (8, 5.5), 1, depth, [2, 3, 7, 10, 11], '4', **dotted)\n",
    "\n",
    "# third block\n",
    "draw_cube(ax, (12, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '5', **solid)\n",
    "draw_cube(ax, (13, 7.5), 1, depth, [1, 2, 3, 6, 9], '8', **solid)\n",
    "draw_cube(ax, (14, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '11', **solid)\n",
    "\n",
    "draw_cube(ax, (12, 6.5), 1, depth, [2, 3, 4], '6', **solid)\n",
    "draw_cube(ax, (13, 6.5), 1, depth, [2, 3], '9', **solid)\n",
    "draw_cube(ax, (14, 6.5), 1, depth, [2, 3, 7, 10], '12', **solid)\n",
    "\n",
    "draw_cube(ax, (12, 5.5), 1, depth, [2, 3, 4], '7', **solid)\n",
    "draw_cube(ax, (13, 5.5), 1, depth, [2, 3], '10', **solid)\n",
    "draw_cube(ax, (14, 5.5), 1, depth, [2, 3, 7, 10], '13', **solid)\n",
    "\n",
    "ax.text(5, 7.0, '+', size=12, ha='center', va='center')\n",
    "ax.text(10.5, 7.0, '=', size=12, ha='center', va='center');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d237f6f3",
   "metadata": {},
   "source": [
    "While broadcasting is very useful, it can sometimes seem confusing.\n",
    "\n",
    "For example, let’s try adding `a -> (3, 2)` and `b -> (3,)`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a383f0f",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array(\n",
    "      [[1, 2],\n",
    "       [4, 5],\n",
    "       [7, 8]])\n",
    "b = np.array([3, 6, 9])\n",
    "\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b3703b7",
   "metadata": {},
   "source": [
    "The `ValueError` tells us that operands could not be broadcast together.\n",
    "\n",
    "Here is a visual representation to show why this broadcasting cannot be executed:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1633c387",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# Draw a figure and axis with no boundary\n",
    "fig = plt.figure(figsize=(3, 1.3), facecolor='w')\n",
    "ax = plt.axes([0, 0, 1, 1], xticks=[], yticks=[], frameon=False)\n",
    "\n",
    "# first block\n",
    "draw_cube(ax, (1, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '1', **solid)\n",
    "draw_cube(ax, (2, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '2', **solid)\n",
    "\n",
    "draw_cube(ax, (1, 6.5), 1, depth, [2, 3, 4], '4', **solid)\n",
    "draw_cube(ax, (2, 6.5), 1, depth, [2, 3, 7, 10], '5', **solid)\n",
    "\n",
    "draw_cube(ax, (1, 5.5), 1, depth, [2, 3, 4], '7', **solid)\n",
    "draw_cube(ax, (2, 5.5), 1, depth, [2, 3, 7, 10], '8', **solid)\n",
    "\n",
    "# second block\n",
    "draw_cube(ax, (6, 7.5), 1, depth, [1, 2, 3, 4, 5, 6, 9], '3', **solid)\n",
    "draw_cube(ax, (7, 7.5), 1, depth, [1, 2, 3, 6, 9], '6', **solid)\n",
    "draw_cube(ax, (8, 7.5), 1, depth, [1, 2, 3, 6, 7, 9, 10], '9', **solid)\n",
    "\n",
    "draw_cube(ax, (6, 6.5), 1, depth, range(2, 13), '3', **dotted)\n",
    "draw_cube(ax, (7, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '6', **dotted)\n",
    "draw_cube(ax, (8, 6.5), 1, depth, [2, 3, 6, 7, 9, 10, 11], '9', **dotted)\n",
    "\n",
    "draw_cube(ax, (6, 5.5), 1, depth, [2, 3, 4, 7, 8, 10, 11, 12], '3', **dotted)\n",
    "draw_cube(ax, (7, 5.5), 1, depth, [2, 3, 7, 10, 11], '6', **dotted)\n",
    "draw_cube(ax, (8, 5.5), 1, depth, [2, 3, 7, 10, 11], '9', **dotted)\n",
    "\n",
    "\n",
    "ax.text(4.5, 7.0, '+', size=12, ha='center', va='center')\n",
    "ax.text(10, 7.0, '=', size=12, ha='center', va='center')\n",
    "ax.text(11, 7.0, '?', size=16, ha='center', va='center');"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91bbb1e1",
   "metadata": {},
   "source": [
    "We can see that NumPy cannot expand the arrays to the same size.\n",
    "\n",
    "It is because, when `b` is expanded from `b -> (3,)` to `b -> (3, 3)`, NumPy cannot match `b` with `a -> (3, 2)`.\n",
    "\n",
    "Things get even trickier when we move to higher dimensions.\n",
    "\n",
    "To help us, we can use the following list of rules:\n",
    "\n",
    "- *Step 1:* When the dimensions of two arrays do not match, NumPy will expand the one with fewer dimensions by adding dimension(s) on the left of the existing dimensions.  \n",
    "  - For example, if `a -> (3, 3)` and `b -> (3,)`, then broadcasting will add a dimension to the left so that `b -> (1, 3)`;  \n",
    "  - If `a -> (2, 2, 2)` and `b -> (2, 2)`, then broadcasting will add a dimension to the left so that `b -> (1, 2, 2)`;  \n",
    "  - If `a -> (3, 2, 2)` and `b -> (2,)`, then broadcasting will add two dimensions to the left so that `b -> (1, 1, 2)` (you can also see this process as going through *Step 1* twice).  \n",
    "- *Step 2:* When the two arrays have the same dimension but different shapes, NumPy will try to expand dimensions where the shape index is 1.  \n",
    "  - For example, if `a -> (1, 3)` and `b -> (3, 1)`, then broadcasting will expand dimensions with shape 1 in both `a` and `b` so that `a -> (3, 3)` and `b -> (3, 3)`;  \n",
    "  - If `a -> (2, 2, 2)` and  `b -> (1, 2, 2)`, then broadcasting will expand the first dimension of `b` so that `b -> (2, 2, 2)`;  \n",
    "  - If `a -> (3, 2, 2)` and `b -> (1, 1, 2)`, then broadcasting will expand `b` on all dimensions with shape 1 so that `b -> (3, 2, 2)`.  \n",
    "\n",
    "\n",
    "Here are code examples for broadcasting higher dimensional arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6d07632d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# a -> (2, 2, 2) and  b -> (1, 2, 2)\n",
    "\n",
    "a = np.array(\n",
    "    [[[1, 2], \n",
    "      [2, 3]], \n",
    "\n",
    "     [[2, 3], \n",
    "      [3, 4]]])\n",
    "print(f'the shape of array a is {a.shape}')\n",
    "\n",
    "b = np.array(\n",
    "    [[1,7],\n",
    "     [7,1]])\n",
    "print(f'the shape of array b is {b.shape}')\n",
    "\n",
    "a + b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2375861",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "# a -> (3, 2, 2) and b -> (2,)\n",
    "\n",
    "a = np.array(\n",
    "    [[[1, 2], \n",
    "      [3, 4]],\n",
    "\n",
    "     [[4, 5], \n",
    "      [6, 7]],\n",
    "\n",
    "     [[7, 8], \n",
    "      [9, 10]]])\n",
    "print(f'the shape of array a is {a.shape}')\n",
    "\n",
    "b = np.array([3, 6])\n",
    "print(f'the shape of array b is {b.shape}')\n",
    "\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f6a318a",
   "metadata": {},
   "source": [
    "- *Step 3:* After Step 1 and 2, if the two arrays still do not match, a `ValueError` will be raised. For example, suppose `a -> (2, 2, 3)` and `b -> (2, 2)`  \n",
    "  - By *Step 1*, `b` will be expanded to `b -> (1, 2, 2)`;  \n",
    "  - By *Step 2*, `b` will be expanded to `b -> (2, 2, 2)`;  \n",
    "  - We can see that they do not match each other after the first two steps. Thus, a `ValueError` will be raised  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f64e9516",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array(\n",
    "    [[[1, 2, 3], \n",
    "      [2, 3, 4]], \n",
    "     \n",
    "     [[2, 3, 4], \n",
    "      [3, 4, 5]]])\n",
    "print(f'the shape of array a is {a.shape}')\n",
    "\n",
    "b = np.array(\n",
    "    [[1,7], \n",
    "     [7,1]])\n",
    "print(f'the shape of array b is {b.shape}')\n",
    "\n",
    "a + b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ca8ce88",
   "metadata": {},
   "source": [
    "## Mutability and Copying Arrays\n",
    "\n",
    "NumPy arrays are mutable data types, like Python lists.\n",
    "\n",
    "In other words, their contents can be altered (mutated) in memory after initialization.\n",
    "\n",
    "We already saw examples above.\n",
    "\n",
    "Here’s another example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5916c6ba",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.array([42, 44])\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1eebc7b6",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a[-1] = 0  # Change last element to 0\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74c55542",
   "metadata": {},
   "source": [
    "Mutability leads to the following behavior (which can be shocking to MATLAB programmers…)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b06cf878",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.random.randn(3)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05a0139d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "b = a\n",
    "b[0] = 0.0\n",
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33040fbb",
   "metadata": {},
   "source": [
    "What’s happened is that we have changed `a` by changing `b`.\n",
    "\n",
    "The name `b` is bound to `a` and becomes just another reference to the\n",
    "array (the Python assignment model is described in more detail [later in the course](https://python-programming.quantecon.org/python_advanced_features.html)).\n",
    "\n",
    "Hence, it has equal rights to make changes to that array.\n",
    "\n",
    "This is in fact the most sensible default behavior!\n",
    "\n",
    "It means that we pass around only pointers to data, rather than making copies.\n",
    "\n",
    "Making copies is expensive in terms of both speed and memory."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0252cbe",
   "metadata": {},
   "source": [
    "### Making Copies\n",
    "\n",
    "It is of course possible to make `b` an independent copy of `a` when required.\n",
    "\n",
    "This can be done using `np.copy`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caae0702",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a = np.random.randn(3)\n",
    "a"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4803d880",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "b = np.copy(a)\n",
    "b"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "485a53d9",
   "metadata": {},
   "source": [
    "Now `b` is an independent copy (called a *deep copy*)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6339a792",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "b[:] = 1\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "175e538a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "a"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "673de7f5",
   "metadata": {},
   "source": [
    "Note that the change to `b` has not affected `a`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bbf4847",
   "metadata": {},
   "source": [
    "## Additional Functionality\n",
    "\n",
    "Let’s look at some other useful things we can do with NumPy."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6570bb62",
   "metadata": {},
   "source": [
    "### Vectorized Functions\n",
    "\n",
    "\n",
    "<a id='index-11'></a>\n",
    "NumPy provides versions of the standard functions `log`, `exp`, `sin`, etc. that act *element-wise* on arrays"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "587ec0c1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.array([1, 2, 3])\n",
    "np.sin(z)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "41aa2e21",
   "metadata": {},
   "source": [
    "This eliminates the need for explicit element-by-element loops such as"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24c5beca",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = len(z)\n",
    "y = np.empty(n)\n",
    "for i in range(n):\n",
    "    y[i] = np.sin(z[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40ad9e6b",
   "metadata": {},
   "source": [
    "Because they act element-wise on arrays, these functions are called *vectorized functions*.\n",
    "\n",
    "In NumPy-speak, they are also called *ufuncs*, which stands for “universal functions”.\n",
    "\n",
    "As we saw above, the usual arithmetic operations (`+`, `*`, etc.) also\n",
    "work element-wise, and combining these with the ufuncs gives a very large set of fast element-wise functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82ae9f78",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cf346202",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "(1 / np.sqrt(2 * np.pi)) * np.exp(- 0.5 * z**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e0087705",
   "metadata": {},
   "source": [
    "Not all user-defined functions will act element-wise.\n",
    "\n",
    "For example, passing the function `f` defined below a NumPy array causes a `ValueError`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0552ee33",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return 1 if x > 0 else 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2ad0aa84",
   "metadata": {},
   "source": [
    "The NumPy function `np.where` provides a vectorized alternative:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6432981",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x = np.random.randn(4)\n",
    "x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e51d09cc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.where(x > 0, 1, 0)  # Insert 1 if x > 0 true, otherwise 0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e6716891",
   "metadata": {},
   "source": [
    "You can also use `np.vectorize` to vectorize a given function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6100d6d8",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "f = np.vectorize(f)\n",
    "f(x)                # Passing the same vector x as in the previous example"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "976c2822",
   "metadata": {},
   "source": [
    "However, this approach doesn’t always obtain the same speed as a more carefully crafted vectorized function."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6c18f6b",
   "metadata": {},
   "source": [
    "### Comparisons\n",
    "\n",
    "\n",
    "<a id='index-12'></a>\n",
    "As a rule, comparisons on arrays are done element-wise"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd31d5e7",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.array([2, 3])\n",
    "y = np.array([2, 3])\n",
    "z == y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "17afaf06",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "y[0] = 5\n",
    "z == y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abc195a0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z != y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72293aa2",
   "metadata": {},
   "source": [
    "The situation is similar for `>`, `<`, `>=` and `<=`.\n",
    "\n",
    "We can also do comparisons against scalars"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57cf7637",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.linspace(0, 10, 5)\n",
    "z"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9c2eca3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z > 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "965d1573",
   "metadata": {},
   "source": [
    "This is particularly useful for *conditional extraction*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "976960b5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "b = z > 3\n",
    "b"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "84224c22",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[b]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1df33354",
   "metadata": {},
   "source": [
    "Of course we can—and frequently do—perform this in one step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "854362b0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z[z > 3]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e300eb",
   "metadata": {},
   "source": [
    "### Sub-packages\n",
    "\n",
    "NumPy provides some additional functionality related to scientific programming\n",
    "through its sub-packages.\n",
    "\n",
    "We’ve already seen how we can generate random variables using np.random"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "821e84e9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "z = np.random.randn(10000)  # Generate standard normals\n",
    "y = np.random.binomial(10, 0.5, size=1000)    # 1,000 draws from Bin(10, 0.5)\n",
    "y.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca8888e1",
   "metadata": {},
   "source": [
    "Another commonly used subpackage is np.linalg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0196497",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "A = np.array([[1, 2], [3, 4]])\n",
    "\n",
    "np.linalg.det(A)           # Compute the determinant"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c7dffa0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.linalg.inv(A)           # Compute the inverse"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "81e19db5",
   "metadata": {},
   "source": [
    "\n",
    "<a id='index-13'></a>\n",
    "\n",
    "<a id='index-14'></a>\n",
    "Much of this functionality is also available in [SciPy](https://www.scipy.org/), a collection of modules that are built on top of NumPy.\n",
    "\n",
    "We’ll cover the SciPy versions in more detail [soon](https://python-programming.quantecon.org/scipy.html).\n",
    "\n",
    "For a comprehensive list of what’s available in NumPy see [this documentation](https://docs.scipy.org/doc/numpy/reference/routines.html)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "96531d21",
   "metadata": {},
   "source": [
    "## Speed Comparisons\n",
    "\n",
    "\n",
    "<a id='index-15'></a>\n",
    "We mentioned in an [previous lecture](https://python-programming.quantecon.org/need_for_speed.html) that NumPy-based vectorization can\n",
    "accelerate scientific applications.\n",
    "\n",
    "In this section we try some speed comparisons to illustrate this fact."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e78d197e",
   "metadata": {},
   "source": [
    "### Vectorization vs Loops\n",
    "\n",
    "Let’s begin with some non-vectorized code, which uses a native Python loop to generate,\n",
    "square and then sum a large number of random variables:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d26b9a68",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "n = 1_000_000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d797176",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "y = 0      # Will accumulate and store sum\n",
    "for i in range(n):\n",
    "    x = random.uniform(0, 1)\n",
    "    y += x**2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc0f8542",
   "metadata": {},
   "source": [
    "The following vectorized code achieves the same thing."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "51cd34d3",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "x = np.random.uniform(0, 1, n)\n",
    "y = np.sum(x**2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01ecd9f4",
   "metadata": {},
   "source": [
    "As you can see, the second code block runs much faster.  Why?\n",
    "\n",
    "The second code block breaks the loop down into three basic operations\n",
    "\n",
    "1. draw `n` uniforms  \n",
    "1. square them  \n",
    "1. sum them  \n",
    "\n",
    "\n",
    "These are sent as batch operators to optimized machine code.\n",
    "\n",
    "Apart from minor overheads associated with sending data back and forth, the result is C or Fortran-like speed.\n",
    "\n",
    "When we run batch operations on arrays like this, we say that the code is *vectorized*.\n",
    "\n",
    "The next section illustrates this point.\n",
    "\n",
    "\n",
    "<a id='ufuncs'></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e9980a8",
   "metadata": {},
   "source": [
    "### Universal Functions\n",
    "\n",
    "\n",
    "<a id='index-16'></a>\n",
    "As discussed above, many functions provided by NumPy are universal functions (ufuncs).\n",
    "\n",
    "By exploiting ufuncs, many operations can be vectorized, leading to faster\n",
    "execution.\n",
    "\n",
    "For example, consider the problem of maximizing a function $ f $ of two\n",
    "variables $ (x,y) $ over the square $ [-a, a] \\times [-a, a] $.\n",
    "\n",
    "For $ f $ and $ a $ let’s choose\n",
    "\n",
    "$$\n",
    "f(x,y) = \\frac{\\cos(x^2 + y^2)}{1 + x^2 + y^2}\n",
    "\\quad \\text{and} \\quad\n",
    "a = 3\n",
    "$$\n",
    "\n",
    "Here’s a plot of $ f $"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7505fdc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def f(x, y):\n",
    "    return np.cos(x**2 + y**2) / (1 + x**2 + y**2)\n",
    "\n",
    "xgrid = np.linspace(-3, 3, 50)\n",
    "ygrid = xgrid\n",
    "x, y = np.meshgrid(xgrid, ygrid)\n",
    "\n",
    "fig = plt.figure(figsize=(10, 8))\n",
    "ax = fig.add_subplot(111, projection='3d')\n",
    "ax.plot_surface(x,\n",
    "                y,\n",
    "                f(x, y),\n",
    "                rstride=2, cstride=2,\n",
    "                cmap=cm.jet,\n",
    "                alpha=0.7,\n",
    "                linewidth=0.25)\n",
    "ax.set_zlim(-0.5, 1.0)\n",
    "ax.set_xlabel('$x$', fontsize=14)\n",
    "ax.set_ylabel('$y$', fontsize=14)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab312db8",
   "metadata": {},
   "source": [
    "To maximize it, we’re going to use a naive grid search:\n",
    "\n",
    "1. Evaluate $ f $ for all $ (x,y) $ in a grid on the square.  \n",
    "1. Return the maximum of observed values.  \n",
    "\n",
    "\n",
    "The grid will be"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32805463",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "grid = np.linspace(-3, 3, 1000)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b1f43f50",
   "metadata": {},
   "source": [
    "Here’s a non-vectorized version that uses Python loops."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2be62009",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "m = -np.inf\n",
    "\n",
    "for x in grid:\n",
    "    for y in grid:\n",
    "        z = f(x, y)\n",
    "        if z > m:\n",
    "            m = z"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0faf783",
   "metadata": {},
   "source": [
    "And here’s a vectorized version"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd6d0cd1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "x, y = np.meshgrid(grid, grid)\n",
    "np.max(f(x, y))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6f644d5",
   "metadata": {},
   "source": [
    "In the vectorized version, all the looping takes place in compiled code.\n",
    "\n",
    "As you can see, the second version is *much* faster."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3a7b70af",
   "metadata": {},
   "source": [
    "## Exercises"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88b4caa0",
   "metadata": {},
   "source": [
    "## Exercise 11.1\n",
    "\n",
    "Consider the polynomial expression\n",
    "\n",
    "\n",
    "<a id='equation-np-polynom'></a>\n",
    "$$\n",
    "p(x) = a_0 + a_1 x + a_2 x^2 + \\cdots a_N x^N = \\sum_{n=0}^N a_n x^n \\tag{11.1}\n",
    "$$\n",
    "\n",
    "[Earlier](https://python-programming.quantecon.org/python_essentials.html#pyess_ex2), you wrote a simple function `p(x, coeff)` to evaluate [(11.1)](#equation-np-polynom) without considering efficiency.\n",
    "\n",
    "Now write a new function that does the same job, but uses NumPy arrays and array operations for its computations, rather than any form of Python loop.\n",
    "\n",
    "(Such functionality is already implemented as `np.poly1d`, but for the sake of the exercise don’t use this class)\n",
    "\n",
    "Use `np.cumprod()`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55ca3915",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 11.1](https://python-programming.quantecon.org/#np_ex1)\n",
    "\n",
    "This code does the job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6a8767b",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "def p(x, coef):\n",
    "    X = np.ones_like(coef)\n",
    "    X[1:] = x\n",
    "    y = np.cumprod(X)   # y = [1, x, x**2,...]\n",
    "    return coef @ y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b76428",
   "metadata": {},
   "source": [
    "Let’s test it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a2aa16a",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "x = 2\n",
    "coef = np.linspace(2, 4, 3)\n",
    "print(coef)\n",
    "print(p(x, coef))\n",
    "# For comparison\n",
    "q = np.poly1d(np.flip(coef))\n",
    "print(q(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14af5bac",
   "metadata": {},
   "source": [
    "## Exercise 11.2\n",
    "\n",
    "Let `q` be a NumPy array of length `n` with `q.sum() == 1`.\n",
    "\n",
    "Suppose that `q` represents a [probability mass function](https://en.wikipedia.org/wiki/Probability_mass_function).\n",
    "\n",
    "We wish to generate a discrete random variable $ x $ such that $ \\mathbb P\\{x = i\\} = q_i $.\n",
    "\n",
    "In other words, `x` takes values in `range(len(q))` and `x = i` with probability `q[i]`.\n",
    "\n",
    "The standard (inverse transform) algorithm is as follows:\n",
    "\n",
    "- Divide the unit interval $ [0, 1] $ into $ n $ subintervals $ I_0, I_1, \\ldots, I_{n-1} $ such that the length of $ I_i $ is $ q_i $.  \n",
    "- Draw a uniform random variable $ U $ on $ [0, 1] $ and return the $ i $ such that $ U \\in I_i $.  \n",
    "\n",
    "\n",
    "The probability of drawing $ i $ is the length of $ I_i $, which is equal to $ q_i $.\n",
    "\n",
    "We can implement the algorithm as follows"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea885c2c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from random import uniform\n",
    "\n",
    "def sample(q):\n",
    "    a = 0.0\n",
    "    U = uniform(0, 1)\n",
    "    for i in range(len(q)):\n",
    "        if a < U <= a + q[i]:\n",
    "            return i\n",
    "        a = a + q[i]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ad492b6",
   "metadata": {},
   "source": [
    "If you can’t see how this works, try thinking through the flow for a simple example, such as `q = [0.25, 0.75]`\n",
    "It helps to sketch the intervals on paper.\n",
    "\n",
    "Your exercise is to speed it up using NumPy, avoiding explicit loops\n",
    "\n",
    "Use `np.searchsorted` and `np.cumsum`\n",
    "\n",
    "If you can, implement the functionality as a class called `DiscreteRV`, where\n",
    "\n",
    "- the data for an instance of the class is the vector of probabilities `q`  \n",
    "- the class has a `draw()` method, which returns one draw according to the algorithm described above  \n",
    "\n",
    "\n",
    "If you can, write the method so that `draw(k)` returns `k` draws from `q`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68800f72",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 11.2](https://python-programming.quantecon.org/#np_ex2)\n",
    "\n",
    "Here’s our first pass at a solution:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "057b7557",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "from numpy import cumsum\n",
    "from numpy.random import uniform\n",
    "\n",
    "class DiscreteRV:\n",
    "    \"\"\"\n",
    "    Generates an array of draws from a discrete random variable with vector of\n",
    "    probabilities given by q.\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, q):\n",
    "        \"\"\"\n",
    "        The argument q is a NumPy array, or array like, nonnegative and sums\n",
    "        to 1\n",
    "        \"\"\"\n",
    "        self.q = q\n",
    "        self.Q = cumsum(q)\n",
    "\n",
    "    def draw(self, k=1):\n",
    "        \"\"\"\n",
    "        Returns k draws from q. For each such draw, the value i is returned\n",
    "        with probability q[i].\n",
    "        \"\"\"\n",
    "        return self.Q.searchsorted(uniform(0, 1, size=k))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85ae29c4",
   "metadata": {},
   "source": [
    "The logic is not obvious, but if you take your time and read it slowly,\n",
    "you will understand.\n",
    "\n",
    "There is a problem here, however.\n",
    "\n",
    "Suppose that `q` is altered after an instance of `discreteRV` is\n",
    "created, for example by"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bfa483c",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "q = (0.1, 0.9)\n",
    "d = DiscreteRV(q)\n",
    "d.q = (0.5, 0.5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a77dffc",
   "metadata": {},
   "source": [
    "The problem is that `Q` does not change accordingly, and `Q` is the\n",
    "data used in the `draw` method.\n",
    "\n",
    "To deal with this, one option is to compute `Q` every time the draw\n",
    "method is called.\n",
    "\n",
    "But this is inefficient relative to computing `Q` once-off.\n",
    "\n",
    "A better option is to use descriptors.\n",
    "\n",
    "A solution from the [quantecon\n",
    "library](https://github.com/QuantEcon/QuantEcon.py/tree/main/quantecon)\n",
    "using descriptors that behaves as we desire can be found\n",
    "[here](https://github.com/QuantEcon/QuantEcon.py/blob/main/quantecon/discrete_rv.py)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37aaf02d",
   "metadata": {},
   "source": [
    "## Exercise 11.3\n",
    "\n",
    "Recall our [earlier discussion](https://python-programming.quantecon.org/python_oop.html#oop_ex1) of the empirical cumulative distribution function.\n",
    "\n",
    "Your task is to\n",
    "\n",
    "1. Make the `__call__` method more efficient using NumPy.  \n",
    "1. Add a method that plots the ECDF over $ [a, b] $, where $ a $ and $ b $ are method parameters.  "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e289b0a9",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 11.3](https://python-programming.quantecon.org/#np_ex3)\n",
    "\n",
    "An example solution is given below.\n",
    "\n",
    "In essence, we’ve just taken [this code](https://github.com/QuantEcon/QuantEcon.py/blob/main/quantecon/ecdf.py)\n",
    "from QuantEcon and added in a plot method"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c3e57cf5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "\"\"\"\n",
    "Modifies ecdf.py from QuantEcon to add in a plot method\n",
    "\n",
    "\"\"\"\n",
    "\n",
    "class ECDF:\n",
    "    \"\"\"\n",
    "    One-dimensional empirical distribution function given a vector of\n",
    "    observations.\n",
    "\n",
    "    Parameters\n",
    "    ----------\n",
    "    observations : array_like\n",
    "        An array of observations\n",
    "\n",
    "    Attributes\n",
    "    ----------\n",
    "    observations : array_like\n",
    "        An array of observations\n",
    "\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, observations):\n",
    "        self.observations = np.asarray(observations)\n",
    "\n",
    "    def __call__(self, x):\n",
    "        \"\"\"\n",
    "        Evaluates the ecdf at x\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        x : scalar(float)\n",
    "            The x at which the ecdf is evaluated\n",
    "\n",
    "        Returns\n",
    "        -------\n",
    "        scalar(float)\n",
    "            Fraction of the sample less than x\n",
    "\n",
    "        \"\"\"\n",
    "        return np.mean(self.observations <= x)\n",
    "\n",
    "    def plot(self, ax, a=None, b=None):\n",
    "        \"\"\"\n",
    "        Plot the ecdf on the interval [a, b].\n",
    "\n",
    "        Parameters\n",
    "        ----------\n",
    "        a : scalar(float), optional(default=None)\n",
    "            Lower endpoint of the plot interval\n",
    "        b : scalar(float), optional(default=None)\n",
    "            Upper endpoint of the plot interval\n",
    "\n",
    "        \"\"\"\n",
    "\n",
    "        # === choose reasonable interval if [a, b] not specified === #\n",
    "        if a is None:\n",
    "            a = self.observations.min() - self.observations.std()\n",
    "        if b is None:\n",
    "            b = self.observations.max() + self.observations.std()\n",
    "\n",
    "        # === generate plot === #\n",
    "        x_vals = np.linspace(a, b, num=100)\n",
    "        f = np.vectorize(self.__call__)\n",
    "        ax.plot(x_vals, f(x_vals))\n",
    "        plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0cfe1f9",
   "metadata": {},
   "source": [
    "Here’s an example of usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f604326",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "X = np.random.randn(1000)\n",
    "F = ECDF(X)\n",
    "F.plot(ax)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c83c4564",
   "metadata": {},
   "source": [
    "## Exercise 11.4\n",
    "\n",
    "Recall that [broadcasting](#broadcasting) in Numpy can help us conduct element-wise operations on arrays with different number of dimensions without using `for` loops.\n",
    "\n",
    "In this exercise, try to use `for` loops to replicate the result of the following broadcasting operations.\n",
    "\n",
    "**Part1**: Try to replicate this simple example using `for` loops and compare your results with the broadcasting operation below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4b32e16",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "x = np.random.randn(4, 4)\n",
    "y = np.random.randn(4)\n",
    "A = x / y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60fbd294",
   "metadata": {},
   "source": [
    "Here is the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "178cb187",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3c4ec49",
   "metadata": {},
   "source": [
    "**Part2**: Move on to replicate the result of the following broadcasting operation. Meanwhile, compare the speeds of broadcasting and the `for` loop you implement.\n",
    "\n",
    "For this part of the exercise you can use the `tic`/`toc` functions from the `quantecon` library to time the execution.\n",
    "\n",
    "Let’s make sure this library is installed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48530fc5",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "!pip install quantecon"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "602c8495",
   "metadata": {},
   "source": [
    "Now we can import the quantecon package."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "549647dc",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "x = np.random.randn(1000, 100, 100)\n",
    "y = np.random.randn(100)\n",
    "\n",
    "qe.tic()\n",
    "B = x / y\n",
    "qe.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1ab71d6",
   "metadata": {},
   "source": [
    "Here is the output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b1fc9356",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(B)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "177545a5",
   "metadata": {},
   "source": [
    "## Solution to[ Exercise 11.4](https://python-programming.quantecon.org/#np_ex4)\n",
    "\n",
    "**Part 1 Solution**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "634cf162",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "x = np.random.randn(4, 4)\n",
    "y = np.random.randn(4)\n",
    "\n",
    "C = np.empty_like(x)\n",
    "n = len(x)\n",
    "for i in range(n):\n",
    "    for j in range(n):\n",
    "        C[i, j] = x[i, j] / y[j]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c07f14d0",
   "metadata": {},
   "source": [
    "Compare the results to check your answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fe3795d",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(C)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "23084b44",
   "metadata": {},
   "source": [
    "You can also use `array_equal()` to check your answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64918a91",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(np.array_equal(A, C))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4eb692cf",
   "metadata": {},
   "source": [
    "**Part 2 Solution**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b120d6a9",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "np.random.seed(123)\n",
    "x = np.random.randn(1000, 100, 100)\n",
    "y = np.random.randn(100)\n",
    "\n",
    "qe.tic()\n",
    "D = np.empty_like(x)\n",
    "d1, d2, d3 = x.shape\n",
    "for i in range(d1):\n",
    "    for j in range(d2):\n",
    "        for k in range(d3):\n",
    "            D[i, j, k] = x[i, j, k] / y[k]\n",
    "qe.toc()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a43912b6",
   "metadata": {},
   "source": [
    "Note that the `for` loop takes much longer than the broadcasting operation.\n",
    "\n",
    "Compare the results to check your answer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ab347bb0",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(D)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a0d66b1",
   "metadata": {
    "hide-output": false
   },
   "outputs": [],
   "source": [
    "print(np.array_equal(B, D))"
   ]
  }
 ],
 "metadata": {
  "date": 1755222216.65885,
  "filename": "numpy.md",
  "kernelspec": {
   "display_name": "Python",
   "language": "python3",
   "name": "python3"
  },
  "title": "NumPy"
 },
 "nbformat": 4,
 "nbformat_minor": 5
}