notebooks/Function_differentiation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "b6f66fb4-ca0e-458c-968e-23bc296429f9",
   "metadata": {},
   "source": [
    "# Function differentiation\n",
    "\n",
    "If you have a Python function `f(x)` that evaluates the mathematical function $f$ at $x$, then we would like to find a function $\\nabla f$ so that we can evaluate $\\nabla f(x)$.\n",
    "\n",
    "Let's look at naive, symbolic and automatic differentiation:\n",
    "\n",
    "- Finite difference\n",
    "- SymPy\n",
    "- JAX\n",
    "\n",
    "💡 There's a TODO at the end of the notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e03fe0f4-66fc-4f28-860c-677b3dc6ae1e",
   "metadata": {},
   "source": [
    "## Examples\n",
    "\n",
    "The examples are from [the Jax docs](https://jax.readthedocs.io/en/latest/automatic-differentiation.html).\n",
    "\n",
    "First, we look at $\\tanh x$, whose derivative is famously convenient to calculate: $1 - \\tanh^2 x$. \n",
    "\n",
    "Second, we look at a polynomial function. Then derivatives of $f(x) = x^3 + 2x^2 - 3x + 1$ can be represented as:\n",
    "\n",
    "$$\n",
    "\\begin{array}{l}\n",
    "f'(x) = 3x^2 + 4x -3\\\\\n",
    "f''(x) = 6x + 4\\\\\n",
    "f'''(x) = 6\\\\\n",
    "f^{iv}(x) = 0\n",
    "\\end{array}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d55fe244-d57c-4d3c-8256-44cf2ffd4e60",
   "metadata": {},
   "source": [
    "## Analytic derivative"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "6eeecbb3-de2b-4690-a9b5-1f12f762e6ee",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.9640275800758169"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import numpy as np\n",
    "\n",
    "x = 2.0\n",
    "\n",
    "f = lambda x: np.tanh(x)\n",
    "f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0fd98dae-bb8c-4f23-a0eb-c02fca703efd",
   "metadata": {},
   "source": [
    "We might know or [look up](https://en.wikipedia.org/wiki/Hyperbolic_functions#Derivatives) that the derivative is given by $1 - \\tanh^2 x$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "4da7ae72-3227-4c8a-a73c-6e8761bb912e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.07065082485316443"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfdx = 1 - f(x)**2\n",
    "dfdx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e161755a-3593-4974-bc4a-c97ed01f9961",
   "metadata": {},
   "source": [
    "## Finite difference\n",
    "\n",
    "### One-sided\n",
    "\n",
    "We can approximate the analytic result with a finite difference, correct to a few decimal places:\n",
    "\n",
    "$$\n",
    "f'(x) \\approx \\frac{f(x + \\epsilon) - f(x)}{\\epsilon}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "30e40efe-bea7-4bcd-bb51-bc875dcf8143",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.07065075680046107"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ϵ = 1e-6\n",
    "dfdx = (f(x + ϵ) - f(x)) / ϵ\n",
    "dfdx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "967d771e-b660-4aae-8db6-4c304390ca0b",
   "metadata": {},
   "source": [
    "### Two-sided\n",
    "\n",
    "We can improve slightly with a symmetric function, but it's still only an estimate, and we cannot get around the floating point imprecision.\n",
    "\n",
    "$$\n",
    "f'(x) \\approx \\frac{f(x + \\epsilon) - f(x - \\epsilon)}{2 \\epsilon}\n",
    "$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "18a9ca9a-14f0-4c71-a1b8-2701d1ad3d48",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.07065082485713248"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ϵ = 1e-6\n",
    "dfdx = (f(x + ϵ) - f(x - ϵ)) / (2 * ϵ)\n",
    "dfdx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58ed4e2a-f108-41fe-94d4-71e6e8373872",
   "metadata": {},
   "source": [
    "## Symbolic differentiation with SymPy\n",
    "\n",
    "Examples inspired by [this Coursera tutorial](https://github.com/greyhatguy007/Mathematics-for-Machine-Learning-and-Data-Science-Specialization-Coursera/blob/main/C2/w1/C2_W1_Lab_1_differentiation_in_python.ipynb) by Luis Serrano."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "31001f99-5695-4706-b9eb-6a73e5c45dd4",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle \\tanh{\\left(x \\right)}$"
      ],
      "text/plain": [
       "tanh(x)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from sympy import tanh, symbols, expand, diff, evalf\n",
    "from sympy.utilities.lambdify import lambdify\n",
    "\n",
    "x = symbols('x')\n",
    "f = tanh(x)\n",
    "f"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b5f49a7-ffe6-4db6-8b8d-37973f4e21bd",
   "metadata": {},
   "source": [
    "We can evaluate for $x = 2$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "edcd8abb-3d50-47f1-9ee6-baebc8bf815e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 0.964027580075817$"
      ],
      "text/plain": [
       "0.964027580075817"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f.evalf(subs={x: 2})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3935eefe-9477-43b2-860d-cc2ed3623cf9",
   "metadata": {},
   "source": [
    "And find the derivative as a symbolic expression:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "bc9bcfb5-e894-4ad4-ad54-72b629b03fac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 1 - \\tanh^{2}{\\left(x \\right)}$"
      ],
      "text/plain": [
       "1 - tanh(x)**2"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfdx = diff(f, x)\n",
    "dfdx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "af8b159b-7317-486e-91b7-52ddc09567eb",
   "metadata": {},
   "source": [
    "And evaluate this:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2f64ae31-cf46-4446-acf5-01a730a3dd58",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 0.0706508248531645$"
      ],
      "text/plain": [
       "0.0706508248531645"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfdx.evalf(subs={x: 2})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "807c652f-5f3c-40de-9c8e-d8ce82632e92",
   "metadata": {},
   "source": [
    "The second derivative simply involves passing the derivative function to `diff()`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ecec19b1-bd9a-434b-8768-cb66edb53efc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 2 \\left(\\tanh^{2}{\\left(x \\right)} - 1\\right) \\tanh{\\left(x \\right)}$"
      ],
      "text/plain": [
       "2*(tanh(x)**2 - 1)*tanh(x)"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d2fdx = diff(f, x, 2)  # Second derivative.\n",
    "d2fdx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b040f1e-db4e-4ffd-a35b-ae61bbdaa020",
   "metadata": {},
   "source": [
    "#### Polynomial example"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4a20f504-eb1b-414f-bd77-710525330b11",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle x^{3} + 2 x^{2} - 3 x + 1$"
      ],
      "text/plain": [
       "x**3 + 2*x**2 - 3*x + 1"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f = x**3 + 2*x**2 - 3*x + 1\n",
    "f"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "029e92f2-4759-4952-9106-a0e17d0fc499",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 11.0$"
      ],
      "text/plain": [
       "11.0000000000000"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f.evalf(subs={x: 2})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "b80fa332-5319-4a3f-8b17-d31c4c38644c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 3 x^{2} + 4 x - 3$"
      ],
      "text/plain": [
       "3*x**2 + 4*x - 3"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfdx = diff(f, x)\n",
    "dfdx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "32b5a138-dc23-478a-a192-b7c93ae24616",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/latex": [
       "$\\displaystyle 17.0$"
      ],
      "text/plain": [
       "17.0000000000000"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfdx.evalf(subs={x: 2})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6b995e6-3b0c-4510-b1c8-668929382ea7",
   "metadata": {},
   "source": [
    "## Automatic differentiation with dual numbers\n",
    "\n",
    "I learned about [dual numbers](https://en.wikipedia.org/wiki/Dual_number) and autodiff from Håvard Berland, who described the method at a company conference, in a version of [part of his PhD defense](https://www.pvv.ntnu.no/~berland/resources/autodiff-triallecture.pdf).\n",
    "\n",
    "Dual numbers are expressions of the form $a + b\\varepsilon$, where $a$ and $b$ are real numbers, and $\\varepsilon$ is a symbol taken to satisfy $\\varepsilon^2 = 0$ with $\\varepsilon\\neq 0$. Evaluating a function with dual numbers produces the derivative automatically:\n",
    "\n",
    "$$ P(a + b\\varepsilon) = P(a) + bP'(a)\\varepsilon $$\n",
    "\n",
    "So we choose $b = 1$.\n",
    "\n",
    "Dual numbers are implemented in some Python libraries, eg [`num-dual`](https://pypi.org/project/num-dual/), but there are others. However, it's not too hard to implement them ourselves."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "3c3f614e-c5bb-43ef-9b10-1cd55189ae9f",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Dual:\n",
    "    def __init__(self, real, dual):\n",
    "        self.real = real\n",
    "        self.dual = dual\n",
    "    def __add__(self, other):\n",
    "        if isinstance(other, Dual):\n",
    "            return Dual(self.real + other.real, self.dual + other.dual)\n",
    "        else:\n",
    "            return Dual(self.real + other, self.dual)\n",
    "    __radd__ = __add__\n",
    "    def __mul__(self, other):\n",
    "        if isinstance(other, Dual):\n",
    "            return Dual(self.real * other.real,\n",
    "                        self.real * other.dual + other.real * self.dual)\n",
    "        else:\n",
    "            return Dual(self.real * other, self.dual * other)\n",
    "    __rmul__ = __mul__\n",
    "    def __neg__(self):\n",
    "        return self.__mul__(-1)\n",
    "    def __sub__(self, other):\n",
    "        return self + -other\n",
    "    def __rsub__(self, other):\n",
    "        return other + -self\n",
    "    def __repr__(self):\n",
    "        return f'Dual({self.real}, {self.dual})'"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a5924ec-bd88-4ba1-8ada-938408688cca",
   "metadata": {},
   "source": [
    "Now we define our function, using multiplication instead of exponentiation (since we have not defined `__pow__()` in our class)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "5d654e0e-62e5-4c71-a8bf-e80d4444dca8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "11"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f = lambda x: x*x*x + 2*x*x - 3*x + 1\n",
    "\n",
    "# Evaluate with x = 2:\n",
    "f(2)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b6d25b60-57e8-45c5-83c8-ff7366333c56",
   "metadata": {},
   "source": [
    "Now for the derivative at $x = 2$. Instead of evaluating on a real number, we evaluate on the dual number, $x = 2 + 1\\epsilon$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "64049bba-faf1-4142-8ec4-5f5cad178776",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Dual(11, 17)"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = Dual(2, 1)\n",
    "\n",
    "f(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f8ad44f-9c7c-47ba-8c07-03d08830dd6f",
   "metadata": {},
   "source": [
    "The dual part, 17, is the derivative."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ca991f39-3841-4e7e-9b04-2752f3790396",
   "metadata": {},
   "source": [
    "## Automatic differentiation with `autograd`\n",
    "\n",
    "[The original `autograd`](https://github.com/HIPS/autograd) from HIPS, now LIPS (at Princeton)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "52045549-0171-42bc-a917-68d84e6c5edb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.419974341614026"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import autograd.numpy as np  # Thinly-wrapped numpy\n",
    "from autograd import grad    # The only autograd function you may ever need\n",
    "\n",
    "def tanh(x):\n",
    "    return (1.0 - np.exp((-2 * x))) / (1.0 + np.exp(-(2 * x)))\n",
    "\n",
    "grad_tanh = grad(tanh)       # Obtain its gradient function\n",
    "grad_tanh(1.0)               # Evaluate the gradient at x = 1.0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8667154c-3494-4591-ae26-225082fdf915",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.41997434264973155"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "(tanh(1.0001) - tanh(0.9999)) / 0.0002  # Compare to finite differences"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00ba04a7-2196-4431-a787-de5f881ee03e",
   "metadata": {},
   "source": [
    "The polynomial example above, recall that the derivative at f(2) was 17..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "b05cd355-bd37-4795-aaae-b8503a3eda74",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "17.0"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = grad(f)\n",
    "\n",
    "df(2.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a123d6a6-bcb6-4a56-b7b6-ffa9de291b4f",
   "metadata": {},
   "source": [
    "## Automatic differentiation with Jax\n",
    "\n",
    "[The Jax docs](https://jax.readthedocs.io/en/latest/automatic-differentiation.html) are really good. All the examples in this notebook were derived from there."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "e5484830-0f04-4082-a33f-b5de1e6a8b7d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.07065082, dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import jax\n",
    "import jax.numpy as jnp\n",
    "from jax import grad\n",
    "\n",
    "grad_tanh = grad(jnp.tanh)\n",
    "\n",
    "x = 2.0\n",
    "grad_tanh(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02bfd32f-9607-4216-a82b-f03a4f2d1a5d",
   "metadata": {},
   "source": [
    "We can get the second-order derivative:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "740e07f0-1375-4c6b-884c-e90e67f06feb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-0.13621868, dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "grad(grad(jnp.tanh))(x)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "58898e5e-7b26-41de-8a13-ee84b3b056d5",
   "metadata": {},
   "source": [
    "For the polynomial:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "53756541-7421-46bf-9b46-46a6344a7dac",
   "metadata": {},
   "outputs": [],
   "source": [
    "f = lambda x: x**3 + 2*x**2 - 3*x + 1\n",
    "\n",
    "dfdx  = jax.grad(f)\n",
    "d2fdx = jax.grad(dfdx)\n",
    "d3fdx = jax.grad(d2fdx)\n",
    "d4fdx = jax.grad(d3fdx)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9d78a6b7-7f36-4b92-ab0a-95ecfd594439",
   "metadata": {},
   "source": [
    "We expect 4.0 here:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "8442e216-4ed8-4e41-a1aa-a046a9576ae0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(4., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfdx(1.0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34002d1c-c4d2-401d-834c-16532a1714ec",
   "metadata": {},
   "source": [
    "## TODO\n",
    "\n",
    "- Compute on arrays\n",
    "- Add plots\n",
    "- Add forward and reverse mode automatic differentiation, eg [see this](https://kenndanielso.github.io/mlrefined/blog_posts/3_Automatic_differentiation/3_4_AD_forward_mode.html)\n",
    "- Also see [Automatic differentiation from scratch](https://e-dorigatti.github.io/math/deep%20learning/2020/04/07/autodiff.html) for a very nice implementation of the computational graph\n",
    "- Add [`tangent`](https://github.com/google/tangent) or whatever replaced it\n",
    "- Add Torch\n",
    "- Add Tensorflow\n",
    "- Have a look at [`yaae`](https://hackmd.io/@machine-learning/blog-post-yaae) (Yet Another Autodiff Engine)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe5d2392-820a-4424-b1b5-7955ef2021b6",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "&copy; Matt Hall 2024 and various original authors linked in text, original content licensed CC BY"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py311",
   "language": "python",
   "name": "py311"
  },
  "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.11.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}