C2/w1/C2_W1_Lab_1_differentiation_in_python.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "reverse-interview",
   "metadata": {},
   "source": [
    "# Differentiation in Python: Symbolic, Numerical and Automatic"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "parental-conclusion",
   "metadata": {},
   "source": [
    "In this lab you explore which tools and libraries are available in Python to compute derivatives. You will perform symbolic differentiation with `SymPy` library, numerical with `NumPy` and automatic with `JAX` (based on `Autograd`). Comparing the speed of calculations, you will investigate the computational efficiency of those three methods."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "looking-barcelona",
   "metadata": {},
   "source": [
    "# Table of Contents\n",
    "- [ 1 - Functions in Python](#1)\n",
    "- [ 2 - Symbolic Differentiation](#2)\n",
    "  - [ 2.1 - Introduction to Symbolic Computation with `SymPy`](#2.1)\n",
    "  - [ 2.2 - Symbolic Differentiation with `SymPy`](#2.2)\n",
    "  - [ 2.3 - Limitations of Symbolic Differentiation](#2.3)\n",
    "- [ 3 - Numerical Differentiation](#3)\n",
    "  - [ 3.1 - Numerical Differentiation with `NumPy`](#3.1)\n",
    "  - [ 3.2 - Limitations of Numerical Differentiation](#3.2)\n",
    "- [ 4 - Automatic Differentiation](#4)\n",
    "  - [ 4.1 - Introduction to `JAX`](#4.1)\n",
    "  - [ 4.2 - Automatic Differentiation with `JAX` ](#4.2)\n",
    "- [ 5 - Computational Efficiency of Symbolic, Numerical and Automatic Differentiation](#5)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "101116ab",
   "metadata": {},
   "source": [
    "<a name='1'></a>\n",
    "## 1 - Functions in Python"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc6140d8",
   "metadata": {},
   "source": [
    "This is just a reminder how to define functions in Python. A simple function $f\\left(x\\right) = x^2$, it can be set up as:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d07a15ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return x**2\n",
    "\n",
    "print(f(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06330bd1",
   "metadata": {},
   "source": [
    "You can easily find the derivative of this function analytically. You can set it up as a separate function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ff4ffb5",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dfdx(x):\n",
    "    return 2*x\n",
    "\n",
    "print(dfdx(3))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8301af3f",
   "metadata": {},
   "source": [
    "Since you have been working with the `NumPy` arrays, you can apply the function to each element of an array:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f5831d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "x_array = np.array([1, 2, 3])\n",
    "\n",
    "print(\"x: \\n\", x_array)\n",
    "print(\"f(x) = x**2: \\n\", f(x_array))\n",
    "print(\"f'(x) = 2x: \\n\", dfdx(x_array))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8428f910",
   "metadata": {},
   "source": [
    "Now you can apply those functions `f` and `dfdx` to an array of a larger size. The following code will plot function and its derivative (you don't have to understand the details of the `plot_f1_and_f2` function at this stage):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c255f4e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# Output of plotting commands is displayed inline within the Jupyter notebook.\n",
    "%matplotlib inline\n",
    "\n",
    "def plot_f1_and_f2(f1, f2=None, x_min=-5, x_max=5, label1=\"f(x)\", label2=\"f'(x)\"):\n",
    "    x = np.linspace(x_min, x_max,100)\n",
    "\n",
    "    # Setting the axes at the centre.\n",
    "    fig = plt.figure()\n",
    "    ax = fig.add_subplot(1, 1, 1)\n",
    "    ax.spines['left'].set_position('center')\n",
    "    ax.spines['bottom'].set_position('zero')\n",
    "    ax.spines['right'].set_color('none')\n",
    "    ax.spines['top'].set_color('none')\n",
    "    ax.xaxis.set_ticks_position('bottom')\n",
    "    ax.yaxis.set_ticks_position('left')\n",
    "\n",
    "    plt.plot(x, f1(x), 'r', label=label1)\n",
    "    if not f2 is None:\n",
    "        # If f2 is an array, it is passed as it is to be plotted as unlinked points.\n",
    "        # If f2 is a function, f2(x) needs to be passed to plot it.        \n",
    "        if isinstance(f2, np.ndarray):\n",
    "            plt.plot(x, f2, 'bo', markersize=3, label=label2,)\n",
    "        else:\n",
    "            plt.plot(x, f2(x), 'b', label=label2)\n",
    "    plt.legend()\n",
    "\n",
    "    plt.show()\n",
    "    \n",
    "plot_f1_and_f2(f, dfdx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ffb711b8",
   "metadata": {},
   "source": [
    "In real life the functions are more complicated and it is not possible to calculate the derivatives analytically every time. Let's explore which tools and libraries are available in Python for the computation of derivatives without manual derivation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d17f76f",
   "metadata": {},
   "source": [
    "<a name='2'></a>\n",
    "## 2 - Symbolic Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "severe-studio",
   "metadata": {},
   "source": [
    "**Symbolic computation** deals with the computation of mathematical objects that are represented exactly, not approximately (e.g. $\\sqrt{2}$ will be written as it is, not as $1.41421356237$). For differentiation it would mean that the output will be somehow similar to if you were computing derivatives by hand using rules (analytically). Thus, symbolic differentiation can produce exact derivatives."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b51ca07",
   "metadata": {},
   "source": [
    "<a name='2.1'></a>\n",
    "### 2.1 - Introduction to Symbolic Computation with `SymPy`\n",
    "\n",
    "Let's explore symbolic differentiation in Python with commonly used `SymPy` library.\n",
    "\n",
    "If you want to compute the approximate decimal value of $\\sqrt{18}$, you could normally do it in the following way:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52d0b0a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "math.sqrt(18)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a5227c24",
   "metadata": {},
   "source": [
    "The output $4.242640687119285$ is an approximate result. You may recall that $\\sqrt{18} = \\sqrt{9 \\cdot 2} = 3\\sqrt{2}$ and see that it is pretty much impossible to deduct it from the approximate result. But with the symbolic computation systems the roots are not approximated with a decimal number but rather only simplified, so the output is exact:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f1747d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This format of module import allows to use the sympy functions without sympy. prefix.\n",
    "from sympy import *\n",
    "\n",
    "# This is actually sympy.sqrt function, but sympy. prefix is omitted.\n",
    "sqrt(18)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5495348",
   "metadata": {},
   "source": [
    "Numerical evaluation of the result is available, and you can set number of the digits to show in the approximated output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "925375b9",
   "metadata": {},
   "outputs": [],
   "source": [
    "N(sqrt(18),8)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb625f44",
   "metadata": {},
   "source": [
    "In `SymPy` variables are defined using **symbols**. In this particular library they need to be predefined (a list of them should be provided). Have a look in the cell below, how the symbolic expression, correspoinding to the mathematical expression $2x^2 - xy$, is defined:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b52721c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# List of symbols.\n",
    "x, y = symbols('x y')\n",
    "# Definition of the expression.\n",
    "expr = 2 * x**2 - x * y\n",
    "expr"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "edffde06",
   "metadata": {},
   "source": [
    "Now you can perform various manipulations with this expression: add or subtract some terms, multiply by other expressions etc., just like if you were doing it by hands:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac575c24",
   "metadata": {},
   "outputs": [],
   "source": [
    "expr_manip = x * (expr + x * y + x**3)\n",
    "expr_manip"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db151265",
   "metadata": {},
   "source": [
    "You can also expand the expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "afb04770",
   "metadata": {},
   "outputs": [],
   "source": [
    "expand(expr_manip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77b3750c",
   "metadata": {},
   "source": [
    "Or factorise it:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0fc456cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "factor(expr_manip)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c885c385",
   "metadata": {},
   "source": [
    "To substitute particular values for the variables in the expression, you can use the following code:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7d239e",
   "metadata": {},
   "outputs": [],
   "source": [
    "expr.evalf(subs={x:-1, y:2})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4acf4536",
   "metadata": {},
   "source": [
    "This can be used to evaluate a function $f\\left(x\\right) = x^2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "622e4335",
   "metadata": {},
   "outputs": [],
   "source": [
    "f_symb = x ** 2\n",
    "f_symb.evalf(subs={x:3})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1346fb9f",
   "metadata": {},
   "source": [
    "You might be wondering now, is it possible to evaluate the symbolic functions for each element of the array? At the beginning of the lab you have defined a `NumPy` array `x_array`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af6562d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(x_array)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f13ae4e0",
   "metadata": {},
   "source": [
    "Now try to evaluate function `f_symb` for each element of the array. You will get an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11df38ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    f_symb(x_array)\n",
    "except TypeError as err:\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f7800fe",
   "metadata": {},
   "source": [
    "It is possible to evaluate the symbolic functions for each element of the array, but you need to make a function `NumPy`-friendly first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d7d8a58",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sympy.utilities.lambdify import lambdify\n",
    "\n",
    "f_symb_numpy = lambdify(x, f_symb, 'numpy')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "631e3e71",
   "metadata": {},
   "source": [
    "The following code should work now:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1ab8f0da",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"x: \\n\", x_array)\n",
    "print(\"f(x) = x**2: \\n\", f_symb_numpy(x_array))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "681062e1",
   "metadata": {},
   "source": [
    "`SymPy` has lots of great functions to manipulate expressions and perform various operations from calculus. More information about them can be found in the official documentation [here](https://docs.sympy.org/)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c6c9d29",
   "metadata": {},
   "source": [
    "<a name='2.2'></a>\n",
    "### 2.2 - Symbolic Differentiation with `SymPy`\n",
    "\n",
    "Let's try to find a derivative of a simple power function using `SymPy`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abb93aeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "diff(x**3,x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd8a81b1",
   "metadata": {},
   "source": [
    "Some standard functions can be used in the expression, and `SymPy` will apply required rules (sum, product, chain) to calculate the derivative:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63ce2cd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_composed = diff(exp(-2*x) + 3*sin(3*x), x)\n",
    "dfdx_composed"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "13407c2b",
   "metadata": {},
   "source": [
    "Now calculate the derivative of the function `f_symb` defined in [2.1](#2.1) and make it `NumPy`-friendly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "597d5879",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_symb = diff(f_symb, x)\n",
    "dfdx_symb_numpy = lambdify(x, dfdx_symb, 'numpy')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f9bee28",
   "metadata": {},
   "source": [
    "Evaluate function `dfdx_symb_numpy` for each element of the `x_array`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b74b4d04",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"x: \\n\", x_array)\n",
    "print(\"f'(x) = 2x: \\n\", dfdx_symb_numpy(x_array))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ada41a99",
   "metadata": {},
   "source": [
    "You can apply symbolically defined functions to the arrays of larger size. The following code will plot function and its derivative, you can see that it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "031a757c",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_f1_and_f2(f_symb_numpy, dfdx_symb_numpy)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de01feee",
   "metadata": {},
   "source": [
    "<a name='2.3'></a>\n",
    "### 2.3 - Limitations of Symbolic Differentiation\n",
    "\n",
    "Symbolic Differentiation seems to be a great tool. But it also has some limitations. Sometimes the output expressions are too complicated and even not possible to evaluate. For example, find the derivative of the function $$\\left|x\\right| = \\begin{cases} x, \\ \\text{if}\\ x > 0\\\\  -x, \\ \\text{if}\\ x < 0 \\\\ 0, \\ \\text{if}\\ x = 0\\end{cases}$$ Analytically, its derivative is:\n",
    "$$\\frac{d}{dx}\\left(\\left|x\\right|\\right) = \\begin{cases} 1, \\ \\text{if}\\ x > 0\\\\  -1, \\ \\text{if}\\ x < 0\\\\\\ \\text{does not exist}, \\ \\text{if}\\ x = 0\\end{cases}$$\n",
    "\n",
    "Have a look the output from the symbolic differentiation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "collect-needle",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_abs = diff(abs(x),x)\n",
    "dfdx_abs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9c3d8e5",
   "metadata": {},
   "source": [
    "Looks complicated, but it would not be a problem if it was possible to evaluate. But check, that for $x=-2$ instead of the derivative value $-1$ it outputs some unevaluated expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d53e7c64",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_abs.evalf(subs={x:-2})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4a9140c",
   "metadata": {},
   "source": [
    "And in the `NumPy` friendly version it also will give an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "644f00ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_abs_numpy = lambdify(x, dfdx_abs,'numpy')\n",
    "\n",
    "try:\n",
    "    dfdx_abs_numpy(np.array([1, -2, 0]))\n",
    "except NameError as err:\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3e1f3c94",
   "metadata": {},
   "source": [
    "In fact, there are problems with the evaluation of the symbolic expressions wherever there is a \"jump\" in the derivative (e.g. function expressions are different for different intervals of $x$), like it happens with $\\frac{d}{dx}\\left(\\left|x\\right|\\right)$. \n",
    "\n",
    "Also, you can see in this example, that you can get a very complicated function as an output of symbolic computation. This is called **expression swell**, which results in unefficiently slow computations. You will see the example of that below after learning other differentiation libraries in Python."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cb85963",
   "metadata": {},
   "source": [
    "<a name='3'></a>\n",
    "## 3 - Numerical Differentiation\n",
    "\n",
    "This method does not take into account the function expression. The only important thing is that the function can be evaluated in the nearby points $x$ and $x+\\Delta x$, where $\\Delta x$ is sufficiently small. Then $\\frac{df}{dx}\\approx\\frac{f\\left(x + \\Delta x\\right) - f\\left(x\\right)}{\\Delta x}$, which can be called a **numerical approximation** of the derivative. \n",
    "\n",
    "Based on that idea there are different approaches for the numerical approximations, which somehow vary in the computation speed and accuracy. However, for all of the methods the results are not accurate - there is a round off error. At this stage there is no need to go into details of various methods, it is enough to investigate one of the numerial differentiation functions, available in `NumPy` package."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1cc2a87e",
   "metadata": {},
   "source": [
    "<a name='3.1'></a>\n",
    "### 3.1 - Numerical Differentiation with `NumPy`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c469b76c",
   "metadata": {},
   "source": [
    "You can call function `np.gradient` to find the derivative of function $f\\left(x\\right) = x^2$ defined above. The first argument is an array of function values, the second defines the spacing $\\Delta x$ for the evaluation. Here pass it as an array of $x$ values, the differences will be calculated automatically. You can find the documentation [here](https://numpy.org/doc/stable/reference/generated/numpy.gradient.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b275519f",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_array_2 = np.linspace(-5, 5, 100)\n",
    "dfdx_numerical = np.gradient(f(x_array_2), x_array_2)\n",
    "\n",
    "plot_f1_and_f2(dfdx_symb_numpy, dfdx_numerical, label1=\"f'(x) exact\", label2=\"f'(x) approximate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6a1d5843",
   "metadata": {},
   "source": [
    "Try to do numerical differentiation for more complicated function:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fa0d7cf",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_composed(x):\n",
    "    return np.exp(-2*x) + 3*np.sin(3*x)\n",
    "\n",
    "plot_f1_and_f2(lambdify(x, dfdx_composed, 'numpy'), np.gradient(f_composed(x_array_2), x_array_2),\n",
    "              label1=\"f'(x) exact\", label2=\"f'(x) approximate\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "826da796",
   "metadata": {},
   "source": [
    "The results are pretty impressive, keeping in mind that it does not matter at all how the function was calculated - only the final values of it!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc60825b",
   "metadata": {},
   "source": [
    "<a name='3.2'></a>\n",
    "### 3.2 - Limitations of Numerical Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dbf76a0",
   "metadata": {},
   "source": [
    "Obviously, the first downside of the numerical differentiation is that it is not exact. However, the accuracy of it is normally enough for machine learning applications. At this stage there is no need to evaluate errors of the numerical differentiation.\n",
    "\n",
    "Another problem is similar to the one which appeared in the symbolic differentiation: it is inaccurate at the points where there are \"jumps\" of the derivative. Let's compare the exact derivative of the absolute value function and with numerical approximation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28bb6a5f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def dfdx_abs(x):\n",
    "    if x > 0:\n",
    "        return 1\n",
    "    else:\n",
    "        if x < 0:\n",
    "            return -1\n",
    "        else:\n",
    "            return None\n",
    "\n",
    "plot_f1_and_f2(np.vectorize(dfdx_abs), np.gradient(abs(x_array_2), x_array_2))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "229fdab5",
   "metadata": {},
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00ccb653",
   "metadata": {},
   "source": [
    "You can see that the results near the \"jump\" are $0.5$ and $-0.5$, while they should be $1$ and $-1$. These cases can give significant errors in the computations.\n",
    "\n",
    "But the biggest problem with the numerical differentiation is slow speed. It requires function evalutation every time.  In machine learning models there are hundreds of parameters and there are hundreds of derivatives to be calculated, performing full function evaluation every time slows down the computation process. You will see the example of it below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2caeb33f",
   "metadata": {},
   "source": [
    "<a name='4'></a>\n",
    "## 4 - Automatic Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eba8f444",
   "metadata": {},
   "source": [
    "**Automatic differentiation** (autodiff) method breaks down the function into common functions ($sin$, $cos$, $log$, power functions, etc.), and constructs the computational graph consisting of the basic functions. Then the chain rule is used to compute the derivative at any node of the graph. It is the most commonly used approach in machine learning applications and neural networks, as the computational graph for the function and its derivatives can be built during the construction of the neural network, saving in future computations.\n",
    "\n",
    "The main disadvantage of it is implementational difficulty. However, nowadays there are libraries that are convenient to use, such as [MyGrad](https://mygrad.readthedocs.io/en/latest/index.html), [Autograd](https://autograd.readthedocs.io/en/latest/) and [JAX](https://jax.readthedocs.io/en/latest/). `Autograd` and `JAX` are the most commonly used in the frameworks to build neural networks. `JAX` brings together `Autograd` functionality for optimization problems, and `XLA` (Accelerated Linear Algebra) compiler for parallel computing.\n",
    "\n",
    "The syntax of `Autograd` and `JAX` are slightly different. It would be overwhelming to cover both at this stage. In this notebook you will be performing automatic differentiation using one of them: `JAX`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20071067",
   "metadata": {},
   "source": [
    "<a name='4.1'></a>\n",
    "### 4.1 - Introduction to `JAX`"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f444d827",
   "metadata": {},
   "source": [
    "To begin with, load the required libraries. From `jax` package you need to load just a couple of functions for now (`grad` and `vmap`). Package `jax.numpy` is a wrapped `NumPy`, which pretty much replaces `NumPy` when `JAX` is used. It can be loaded as `np` as if it was an original `NumPy` in most of the cases. However, in this notebook you'll upload it as `jnp` to distinguish them for now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85d818ea",
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import grad, vmap\n",
    "import jax.numpy as jnp"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d42ede8e",
   "metadata": {},
   "source": [
    "Create a new `jnp` array and check its type."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8856647c",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_array_jnp = jnp.array([1., 2., 3.])\n",
    "\n",
    "print(\"Type of NumPy array:\", type(x_array))\n",
    "print(\"Type of JAX NumPy array:\", type(x_array_jnp))\n",
    "# Please ignore the warning message if it appears."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "730a2dd3",
   "metadata": {},
   "source": [
    "The same array can be created just converting previously defined `x_array = np.array([1, 2, 3])`, although in some cases `JAX` does not operate with integers, thus the values need to be converted to floats. You will see an example of it below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3008671b",
   "metadata": {},
   "outputs": [],
   "source": [
    "x_array_jnp = jnp.array(x_array.astype('float32'))\n",
    "print(\"JAX NumPy array:\", x_array_jnp)\n",
    "print(\"Type of JAX NumPy array:\", type(x_array_jnp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f81ce077",
   "metadata": {},
   "source": [
    "Note, that `jnp` array has a specific type `jaxlib.xla_extension.DeviceArray`. In most of the cases the same operators and functions are applicable to them as in the original `NumPy`, for example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "742003ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(x_array_jnp * 2)\n",
    "print(x_array_jnp[2])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c7ef8a4",
   "metadata": {},
   "source": [
    "But sometimes working with `jnp` arrays the approach needs to be changed. In the following code, trying to assign a new value to one of the elements, you will get an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3fc00cab",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    x_array_jnp[2] = 4.0\n",
    "except TypeError as err:\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cf9e29fe",
   "metadata": {},
   "source": [
    "To assign a new value to an element in the `jnp` array you need to apply functions `.at[i]`, stating which element to update, and `.set(value)` to set a new value. These functions also operate **out-of-place**, the updated array is returned as a new array and the original array is not modified by the update."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ffc53ad2",
   "metadata": {},
   "outputs": [],
   "source": [
    "y_array_jnp = x_array_jnp.at[2].set(4.0)\n",
    "print(y_array_jnp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05a07ce0",
   "metadata": {},
   "source": [
    "Although, some of the `JAX` functions will work with arrays defined with `np` and `jnp`. In the following code you will get the same result in both lines:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b80429d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(jnp.log(x_array))\n",
    "print(jnp.log(x_array_jnp))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89397092",
   "metadata": {},
   "source": [
    "This is probably confusing - which `NumPy` to use then? Usually when `JAX` is used, only `jax.numpy` gets imported as `np`, and used instead of the original one."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20f12b94",
   "metadata": {},
   "source": [
    " <a name='4.2'></a>\n",
    "### 4.2 - Automatic Differentiation with `JAX` "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9cd26792",
   "metadata": {},
   "source": [
    "Time to do automatic differentiation with `JAX`. The following code will calculate the derivative of the previously defined function $f\\left(x\\right) = x^2$ at the point $x = 3$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "070e417a",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Function value at x = 3:\", f(3.0))\n",
    "print(\"Derivative value at x = 3:\",grad(f)(3.0))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3514bda9",
   "metadata": {},
   "source": [
    "Very easy, right? Keep in mind, please, that this cannot be done using integers. The following code will output an error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a50295a3",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    grad(f)(3)\n",
    "except TypeError as err:\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "872bbbc6",
   "metadata": {},
   "source": [
    "Try to apply the `grad` function to an array, calculating the derivative for each of its elements: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caf0e431",
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    grad(f)(x_array_jnp)\n",
    "except TypeError as err:\n",
    "    print(err)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9452ebc2",
   "metadata": {},
   "source": [
    "There is some broadcasting issue there. You don't need to get into more details of this at this stage, function `vmap` can be used here to solve the problem.\n",
    "\n",
    "*Note*: Broadcasting is covered in the Course 1 of this Specialization \"Linear Algebra\". You can also review it in the documentation [here](https://numpy.org/doc/stable/user/basics.broadcasting.html#:~:text=The%20term%20broadcasting%20describes%20how,that%20they%20have%20compatible%20shapes.)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9b28641",
   "metadata": {},
   "outputs": [],
   "source": [
    "dfdx_jax_vmap = vmap(grad(f))(x_array_jnp)\n",
    "print(dfdx_jax_vmap)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "933e382f",
   "metadata": {},
   "source": [
    "Great, now `vmap(grad(f))` can be used to calculate the derivative of function `f` for arrays of larger size and you can plot the output:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da0a1262",
   "metadata": {},
   "outputs": [],
   "source": [
    "plot_f1_and_f2(f, vmap(grad(f)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4162d5e5",
   "metadata": {},
   "source": [
    "In the following code you can comment/uncomment lines to visualize the common derivatives. All of them are found using `JAX` automatic differentiation. The results look pretty good!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f68b4c0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def g(x):\n",
    "#     return x**3\n",
    "#     return 2*x**3 - 3*x**2 + 5\n",
    "#     return 1/x\n",
    "#     return jnp.exp(x)\n",
    "#     return jnp.log(x)\n",
    "#     return jnp.sin(x)\n",
    "#     return jnp.cos(x)\n",
    "    return jnp.abs(x)\n",
    "#     return jnp.abs(x)+jnp.sin(x)*jnp.cos(x)\n",
    "\n",
    "plot_f1_and_f2(g, vmap(grad(g)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a58ee858",
   "metadata": {},
   "source": [
    "<a name='5'></a>\n",
    "## 5 - Computational Efficiency of Symbolic, Numerical and Automatic Differentiation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2211158e",
   "metadata": {},
   "source": [
    "In sections [2.3](#2.3) and [3.2](#3.2) low computational efficiency of symbolic and numerical differentiation was discussed. Now it is time to compare speed of calculations for each of three approaches. Try to find the derivative of the same simple function $f\\left(x\\right) = x^2$ multiple times, evaluating it for an array of a larger size, compare the results and time used:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36c42dac",
   "metadata": {},
   "outputs": [],
   "source": [
    "import timeit, time\n",
    "\n",
    "x_array_large = np.linspace(-5, 5, 1000000)\n",
    "\n",
    "tic_symb = time.time()\n",
    "res_symb = lambdify(x, diff(f(x),x),'numpy')(x_array_large)\n",
    "toc_symb = time.time()\n",
    "time_symb = 1000 * (toc_symb - tic_symb)  # Time in ms.\n",
    "\n",
    "tic_numerical = time.time()\n",
    "res_numerical = np.gradient(f(x_array_large),x_array_large)\n",
    "toc_numerical = time.time()\n",
    "time_numerical = 1000 * (toc_numerical - tic_numerical)\n",
    "\n",
    "tic_jax = time.time()\n",
    "res_jax = vmap(grad(f))(jnp.array(x_array_large.astype('float32')))\n",
    "toc_jax = time.time()\n",
    "time_jax = 1000 * (toc_jax - tic_jax)\n",
    "\n",
    "print(f\"Results\\nSymbolic Differentiation:\\n{res_symb}\\n\" + \n",
    "      f\"Numerical Differentiation:\\n{res_numerical}\\n\" + \n",
    "      f\"Automatic Differentiation:\\n{res_jax}\")\n",
    "\n",
    "print(f\"\\n\\nTime\\nSymbolic Differentiation:\\n{time_symb} ms\\n\" + \n",
    "      f\"Numerical Differentiation:\\n{time_numerical} ms\\n\" + \n",
    "      f\"Automatic Differentiation:\\n{time_jax} ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "493e5457",
   "metadata": {},
   "source": [
    "The results are pretty much the same, but the time used is different. Numerical approach is obviously inefficient when differentiation needs to be performed many times, which happens a lot training machine learning models. Symbolic and automatic approach seem to be performing similarly for this simple example. But if the function becomes a little bit more complicated, symbolic computation will experiance significant expression swell and the calculations will slow down.\n",
    "\n",
    "*Note*: Sometimes the execution time results may vary slightly, especially for automatic differentiation. You can run the code above a few time to see different outputs. That does not influence the conclusion that numerical differentiation is slower. `timeit` module can be used more efficiently to evaluate execution time of the codes, but that would unnecessary overcomplicate the codes here.\n",
    "\n",
    "Try to define some polynomial function, which should not be that hard to differentiate, and compare the computation time for its differentiation symbolically and automatically:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13047a93",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f_polynomial_simple(x):\n",
    "    return 2*x**3 - 3*x**2 + 5\n",
    "\n",
    "def f_polynomial(x):\n",
    "    for i in range(3):\n",
    "        x = f_polynomial_simple(x)\n",
    "    return x\n",
    "\n",
    "tic_polynomial_symb = time.time()\n",
    "res_polynomial_symb = lambdify(x, diff(f_polynomial(x),x),'numpy')(x_array_large)\n",
    "toc_polynomial_symb = time.time()\n",
    "time_polynomial_symb = 1000 * (toc_polynomial_symb - tic_polynomial_symb)\n",
    "\n",
    "tic_polynomial_jax = time.time()\n",
    "res_polynomial_jax = vmap(grad(f_polynomial))(jnp.array(x_array_large.astype('float32')))\n",
    "toc_polynomial_jax = time.time()\n",
    "time_polynomial_jax = 1000 * (toc_polynomial_jax - tic_polynomial_jax)\n",
    "\n",
    "print(f\"Results\\nSymbolic Differentiation:\\n{res_polynomial_symb}\\n\" + \n",
    "      f\"Automatic Differentiation:\\n{res_polynomial_jax}\")\n",
    "\n",
    "print(f\"\\n\\nTime\\nSymbolic Differentiation:\\n{time_polynomial_symb} ms\\n\" +  \n",
    "      f\"Automatic Differentiation:\\n{time_polynomial_jax} ms\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "231c9da0",
   "metadata": {},
   "source": [
    "Again, the results are similar, but automatic differentiation is times faster. \n",
    "\n",
    "With the increase of function computation graph, the efficiency of automatic differentiation compared to other methods raises, because autodiff method uses chain rule!\n",
    "\n",
    "Congratulations! Now you are equiped with Python tools to perform differentiation."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93dfd229",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}