notebooks/4. AutoDiff.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Autodiff:  <br> Calculus  from another angle \n",
    "(and the special role played by Julia's multiple dispatch and compiler technology)\n",
    "\n",
    "(Alan Edelman, 2017)\n",
    "\n",
    "\n",
    "  The first time I heard about automatic differentiation, it was easy for me to imagine what it was.  I was wrong.  In my head, I thought it was straightforward symbolic differentiation applied to code.  I kind of imagined it was like executing Mathematica or Maple, or even just automatically doing what I learned to do in my calculus class. \n",
    "  <img src=\"https://math24.net/images/table-trig-derivatives.svg\" width=\"290\">\n",
    "  .... and anyway if it was not that, then it must be finite differences, like one learns in a numerical computing class.\n",
    "  \n",
    "<img src=\"https://upload.wikimedia.org/wikipedia/commons/thumb/9/90/Finite_difference_method.svg/614px-Finite_difference_method.svg.png\" width=\"350\">\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Babylonian sqrt\n",
    "\n",
    "I would like to use a simple example, computation of sqrt(x), where for me how autodiff works came as both a mathematical surprise, and a computing wonder.  The example is  the Babylonian algorithm, known to man for millenia, to compute sqrt(x):  \n",
    "\n",
    "\n",
    " > Repeat $ t \\leftarrow  (t+x/t) / 2 $ until $t$ converges to $\\sqrt{x}$.\n",
    " \n",
    " Each iteration has one add and two divides. For illustration purposes, 10 iterations suffice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function Babylonian(x; N = 10) \n",
    "    t = (1+x)/2\n",
    "    for i = 2:N; t=(t + x/t)/2  end    \n",
    "    t\n",
    "end  "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Check that it works:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "α = π\n",
    "Babylonian(α), √α    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=2; Babylonian(x),√x  # Type \\sqrt+<tab> to get the symbol"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Plots\n",
    "plotly()\n",
    "#gr()\n",
    "#pyplot()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = 0:.01:49\n",
    "\n",
    "plot([x->Babylonian(x,N=i) for i=1:5],i,label=[\"Iteration $j\" for i=1:1,j=1:5])\n",
    "\n",
    "plot!(sqrt,i,c=\"black\",label=\"sqrt\",\n",
    "      title = \"Those Babylonians really knew how to √\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ...and now the derivative, almost by magic"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "struct D <: Number  # D is a function-derivative pair\n",
    "    f::Tuple{Float64,Float64}\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D((1.2,3.14))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D(x,y) = D((x,y)) # teaching point, overload D to take in either a tupe, or two arguments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D(1.2,3.14)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D(1,2) / D(3,4)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Sum Rule: (x+y)' = x' + y' <br>\n",
    "Quotient Rule: (x/y)' = (yx'-xy') / y^2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import Base: +, /, convert, promote_rule\n",
    "+(x::D, y::D) = D(x.f .+ y.f)\n",
    "/(x::D, y::D) = D((x.f[1]/y.f[1], (y.f[1]*x.f[2] - x.f[1]*y.f[2])/y.f[1]^2))\n",
    "convert(::Type{D}, x::Real) = D((x,zero(x)))\n",
    "promote_rule(::Type{D}, ::Type{<:Number}) = D"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D(2,3) / D(3,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=49; Babylonian(D(x,1)), (√x,.5/√x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x=π; Babylonian(D((x,1))), (√x,.5/√x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Babylonian(  D(100,1)   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "i = .2:.01:49\n",
    "plot([x->Babylonian(D((x,1.0)),N=i).f[2] for i=1:5],i)\n",
    "plot!(x->.5/√x,i,c=\"black\",label=\"d(sqrt(x))/dx\",\n",
    "    title = \" Babylonians Differentiated\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## It just works!\n",
    "\n",
    "How does it work?  We will explain in a moment.  Right now marvel that it does.  Note we did not\n",
    "import any autodiff package.  Everything is just basic vanilla Julia."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The assembler\n",
    "\n",
    "Most folks don't read assembler, but one can see that it is short."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# @inline function Babylonian(x; N = 10) \n",
    "#     t = (1+x)/2\n",
    "#     for i = 2:N; t=(t + x/t)/2  end    \n",
    "#     t\n",
    "# end  \n",
    "@code_native(Babylonian(D((2,1))))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Symbolically\n",
    "\n",
    "We haven't yet explained how it works, but it may be of some value to understand that the below is mathematically\n",
    "equivalent, though not what the computation is doing.\n",
    "\n",
    "Notice in the below that Babylonian works on SymPy symbols."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Pkg.add(\"SymPy\")\n",
    "using SymPy                    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = symbols(\"x\")\n",
    "display(\"Iterations as a function of x\")\n",
    "for k = 1:5\n",
    " display( simplify(Babylonian(x,N=k)))\n",
    "end\n",
    "\n",
    "display(\"Derivatives as a function of x\")\n",
    "for k = 1:5\n",
    " display(simplify(diff(simplify(Babylonian(x,N=k)),x)))\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How autodiff is getting the answer\n",
    "Let us by hand take the \"derivative\" of the Babylonian iteration with respect to x. Specifically t′=dt/dx"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function dBabylonian(x; N = 10) \n",
    "    t = (1+x)/2\n",
    "    t′ = 1/2\n",
    "    for i = 1:N;  \n",
    "        t = (t+x/t)/2; \n",
    "        t′= (t′+(t-x*t′)/t^2)/2; \n",
    "    end    \n",
    "    t′\n",
    "\n",
    "end  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = π; dBabylonian(x), .5/√x"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What just happened?  Answer: We created an iteration by hand for t′ given our iteration for t. Then we ran the iteration alongside the iteration for t."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Babylonian(D((x,1)))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "How did this work?  It created the same derivative iteration, using very general rules that are set once and need not be written by hand."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Important:: The derivative is substituted before the JIT compiler, and thus efficient compiled code is executed."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Dual Number Notation\n",
    "\n",
    "Instead of D(a,b) we can write a + b ϵ, where ϵ satisfies ϵ^2=0.  (Some people like to recall imaginary numbers where an i is introduced with i^2=-1.) \n",
    "\n",
    "Others like to think of how engineers just drop the O(ϵ^2) terms.\n",
    "\n",
    "The four rules are\n",
    "\n",
    "$ (a+b\\epsilon) \\pm (c+d\\epsilon) = (a \\pm c) +  (b \\pm d)\\epsilon$\n",
    "\n",
    "$ (a+b\\epsilon) * (c+d\\epsilon) = (ac) + (bc+ad)\\epsilon$\n",
    "\n",
    "$ (a+b\\epsilon) / (c+d\\epsilon) = (a/c) + (bc-ad)/c^2 \\epsilon $\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@syms a b c d ϵ"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "series((a+b*ϵ)/(c+d*ϵ),ϵ,0,2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Base.show(io::IO,x::D) = print(io,x.f[1],\" + \",x.f[2],\" ϵ\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D(4,5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Add the last two rules\n",
    "import Base: -,*\n",
    "-(x::D, y:: D) = D(x.f .- y.f)\n",
    "*(x::D, y::D) = D((x.f[1]*y.f[1], (x.f[2]*y.f[1] + x.f[1]*y.f[2])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D((1,0))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "D((2,1)) ^2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ϵ = D((0,1))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ϵ * ϵ "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ϵ^2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "1/(1+ϵ) \n",
    "# Exact power series:  1-ϵ+ϵ²-ϵ³-..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(1+ϵ)^5  ## Note this just works (we didn't train powers)!!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "(1+ϵ)^(-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generalization to arbitrary roots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "function nthroot(x, n=2; t=1, N = 10) \n",
    "    for i = 1:N;   t += (x/t^(n-1)-t)/n; end   \n",
    "    t\n",
    "end  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nthroot(2,3), ∛2 # take a cube root"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nthroot(2+ϵ,3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nthroot(7,12), 7^(1/12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x = 2.0\n",
    "nthroot( x+ϵ,3), ∛x, 1/x^(2/3)/3"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Forward Diff\n",
    "Now that you understand it, you can use the official package"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using Pkg\n",
    "Pkg.add(\"ForwardDiff\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "using ForwardDiff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ForwardDiff.derivative(sqrt, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ForwardDiff.derivative(Babylonian, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@which ForwardDiff.derivative(sqrt, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Julia 1.9.3",
   "language": "julia",
   "name": "julia-1.9"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}