psets/pset2sol.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "25876c85",
   "metadata": {},
   "source": [
    "# Matrix Calculus Pset 2 Solutions\n",
    "\n",
    "These is the computational notebook accompanying the pset 2 solutions for *Matrix Calculus*, MIT IAP 2023."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7ec6213c",
   "metadata": {},
   "source": [
    "# Problem 1(c)\n",
    "\n",
    "## f(p)\n",
    "\n",
    "First, let's translate $f(p)$ to Julia.  Recall that this is defined by $f(p) = \\left(c^{T}A(p)^{-1}b\\right)^{2}$ where\n",
    "$$\n",
    "A(p)=\\left(\\begin{array}{ccccc}\n",
    "a_{1} & p_{1}\\\\\n",
    "p_{1} & a_{2} & p_{2}\\\\\n",
    " & p_{2} & \\ddots & \\ddots\\\\\n",
    " &  & \\ddots & a_{n-1} & p_{n-1}\\\\\n",
    " &  &  & p_{n-1} & a_{n}\n",
    "\\end{array}\\right),\n",
    "$$\n",
    "for some constant vectors $a,b,c$.\n",
    "\n",
    "Let's make $f$ for $n=5$, for example.  We'll use Julia's handy [`LinearAlgebra.SymTridiagonal` type](https://docs.julialang.org/en/v1/stdlib/LinearAlgebra/#LinearAlgebra.SymTridiagonal) to make $A(p)$ first:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "4942e6d0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5×5 SymTridiagonal{Float64, Vector{Float64}}:\n",
       " -9.0       0.416908   ⋅          ⋅          ⋅ \n",
       "  0.416908  1.0       0.292966    ⋅          ⋅ \n",
       "   ⋅        0.292966  0.0        0.310202    ⋅ \n",
       "   ⋅         ⋅        0.310202  -6.0        0.485178\n",
       "   ⋅         ⋅         ⋅         0.485178  -1.0"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n = 5\n",
    "\n",
    "# some random vectors\n",
    "a = rand(-9:9, n)\n",
    "b = rand(-9:9, n)\n",
    "c = rand(-9:9, n)\n",
    "\n",
    "using LinearAlgebra # for SymTridiagonal\n",
    "\n",
    "A(p) = SymTridiagonal(promote(a, p)...)\n",
    "\n",
    "p = rand(n-1) # for example\n",
    "A(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b6063b4",
   "metadata": {},
   "source": [
    "Now, $f(p)$ is simple:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b3ecdf88",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13895.37546070443"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f(p) = (c' * (A(p) \\ b))^2\n",
    "\n",
    "f(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30b700b1",
   "metadata": {},
   "source": [
    "Note that `A(p) \\ b` in Julia automatically uses a specialized $\\Theta(n)$ tridiagonal solver since `A` is the `SymTridiagonal` type."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2304d51d",
   "metadata": {},
   "source": [
    "##  adjoint/reverse-mode ∇f\n",
    "\n",
    "Now, let's implement our \"adjoint-method\" or \"reverse-mode\" algorithm for $\\nabla f$:\n",
    "\n",
    "1. Solve the forward problem for $x = A(p)^{-1} b$\n",
    "2. Solve the adjoint problem $v = A^{-1} [-2(c^T x) c]$\n",
    "3. $\\partial f/\\partial p_k = v_k x_{k+1} + v_{k+1} x_k$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "74e821ce",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "∇f (generic function with 1 method)"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function ∇f(p)\n",
    "    M = A(p)\n",
    "    x = M \\ b\n",
    "    v = M \\ (-2(c'x)c)\n",
    "    return [ v[k]*x[k+1] + v[k+1]*x[k] for k = 1:length(p) ]\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "1293c81b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4-element Vector{Float64}:\n",
       "    8672.424124240617\n",
       " -234089.45857983222\n",
       "   69396.70127620807\n",
       "    8926.845054799993"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∇f(p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2e1cbd39",
   "metadata": {},
   "source": [
    "## Finite-difference check\n",
    "\n",
    "Now let's do a finite-difference check that\n",
    "$$\n",
    "f(p+\\delta) - f(p) \\approx \\nabla f^T \\delta\n",
    "$$\n",
    "for a small random $\\delta$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "41aff3c0",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00016509416491317097"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "δ = randn(n-1) * 1e-8\n",
    "\n",
    "df = f(p+δ) - f(p)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4633e5e8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.00016509415974040186"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∇f(p)'δ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae992ae7",
   "metadata": {},
   "source": [
    "Looks good!  They match to quite a few digits!\n",
    "\n",
    "We can be more quantitative by computing the relative error:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "6007c89f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.1332235629796853e-8"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(a,b) = norm(a - b) / norm(b)\n",
    "\n",
    "relerr(df, ∇f(p)'δ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1c20440",
   "metadata": {},
   "source": [
    "So, the relative error is about $10^{-8}$, pretty reasonable for $\\delta$ on that order."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3151f23a",
   "metadata": {},
   "source": [
    "## AD check\n",
    "\n",
    "For fun, let's also try automatic differentiation via ForwardDiff.jl or Zygote.jl (this was **not** required):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "97b98ce3",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ForwardDiff, Zygote"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "13c98f9f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.4586950654248017e-16"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(ForwardDiff.gradient(f, p), ∇f(p))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e1e2acf4",
   "metadata": {},
   "source": [
    "Great!  They match to almost [machine precision](https://en.wikipedia.org/wiki/Machine_epsilon)!\n",
    "\n",
    "Let's try Zygote:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "d494c8bb",
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "Need an adjoint for constructor SymTridiagonal{Float64, Vector{Float64}}. Gradient is of type Matrix{Float64}",
     "output_type": "error",
     "traceback": [
      "Need an adjoint for constructor SymTridiagonal{Float64, Vector{Float64}}. Gradient is of type Matrix{Float64}",
      "",
      "Stacktrace:",
      "  [1] error(s::String)",
      "    @ Base ./error.jl:35",
      "  [2] (::Zygote.Jnew{SymTridiagonal{Float64, Vector{Float64}}, Nothing, false})(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/lib/lib.jl:330",
      "  [3] (::Zygote.var\"#2100#back#226\"{Zygote.Jnew{SymTridiagonal{Float64, Vector{Float64}}, Nothing, false}})(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67",
      "  [4] Pullback",
      "    @ /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/tridiag.jl:14 [inlined]",
      "  [5] (::typeof(∂(SymTridiagonal{Float64, Vector{Float64}})))(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface2.jl:0",
      "  [6] Pullback",
      "    @ /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/tridiag.jl:71 [inlined]",
      "  [7] Pullback",
      "    @ /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/tridiag.jl:70 [inlined]",
      "  [8] #208",
      "    @ ~/.julia/packages/Zygote/AS0Go/src/lib/lib.jl:206 [inlined]",
      "  [9] #2066#back",
      "    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]",
      " [10] Pullback",
      "    @ ./In[1]:10 [inlined]",
      " [11] (::typeof(∂(A)))(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface2.jl:0",
      " [12] Pullback",
      "    @ ./In[2]:1 [inlined]",
      " [13] (::Zygote.var\"#60#61\"{typeof(∂(f))})(Δ::Float64)",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface.jl:45",
      " [14] gradient(f::Function, args::Vector{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface.jl:97",
      " [15] top-level scope",
      "    @ In[10]:1"
     ]
    }
   ],
   "source": [
    "relerr(Zygote.gradient(f, p), ∇f(p))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5c36a37",
   "metadata": {},
   "source": [
    "Whoops, Zygote doesn't know how to differentiate through the `SymTridiagonal` function.  Grr."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2cd2ba2",
   "metadata": {},
   "source": [
    "# Problem 4b\n",
    "\n",
    "Our test matrix $M$ is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9f7af454",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×3 Matrix{Int64}:\n",
       " 0  1  4\n",
       " 1  0  1\n",
       " 4  1  0"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "M = [0 1 4; 1 0 1; 4 1 0]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac64fbc3",
   "metadata": {},
   "source": [
    "## (i) By finite differences\n",
    "\n",
    "Here, we want the \"matrix Jacobian determinant\" for a matrix-valued function $f(M)$, defined by $\\det J$ where $\\text{vec}(df) = J \\text{vec}(M)$.\n",
    "\n",
    "Let's do it by (forward) finite differences, column by column in $J$.\n",
    "\n",
    "One handy trick is that if you have an $m \\times n$ matrix `M` in Julia, then you can index it by row/column as `M[i,j]` **or** you can index it \"linearly\" as `M[i]` for $i = 1,\\ldots,mn$ as if it were an $mn$-component vector, exactly corresponding to $\\text{vec}(M)$ (since this is, in fact, how Julia stores the matrix in your computer's \"1d\" memory).  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "baee4911",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fd_jacobian (generic function with 2 methods)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function fd_jacobian(f, M, δ=1e-8)\n",
    "    fM = f(M)\n",
    "    Mplusδ = float.(M) # make a floating-point copy of M\n",
    "    J = zeros(length(fM), length(M))\n",
    "    for j = 1:length(M)\n",
    "        Mplusδ[j] += δ\n",
    "        fMplusδ = f(Mplusδ)\n",
    "        for i = 1:length(fM)\n",
    "            J[i,j] = (fMplusδ[i] - fM[i]) / δ\n",
    "        end\n",
    "        Mplusδ[j] = M[j] # reset for next finite difference\n",
    "    end\n",
    "    return J\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "fd447a6a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " 22.9849    7.49817  18.369     7.49817  …  18.369     5.4667   13.7714\n",
       "  7.49818  11.3611    7.49818   2.1476       5.46671  10.2735    5.46671\n",
       " 18.369     7.49817  22.9849    5.4667      13.7714    5.4667   18.369\n",
       "  7.49818   2.1476    5.46671  11.3611       7.49818   2.14759   5.46671\n",
       "  2.1476    4.37449   2.1476    4.37449      2.1476    4.37449   2.1476\n",
       "  5.46671   2.1476    7.49818  10.2735   …   5.46671   2.1476    7.49818\n",
       " 18.369     5.4667   13.7714    7.49817     22.9849    7.49817  18.369\n",
       "  5.46671  10.2735    5.46671   2.1476       7.49818  11.3611    7.49818\n",
       " 13.7714    5.46671  18.369     5.46671     18.369     7.49817  22.9849"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "fd_jacobian(exp, M)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2897d982",
   "metadata": {},
   "source": [
    "Now, let's use that to compute the Jacobian determinant for $f \\in {\\exp, \\sin}$.\n",
    "\n",
    "Note that, in Julia, you can compute the matrix exponential or sine simply by `exp(M)` or `sin(M)` respectively — these are *not* the elementwise functions.   (In Scipy or Matlab you would need `expm` and `sinm`.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "ee176095",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "939.0318094614613"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jexp_fd = det(fd_jacobian(exp, M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "b7b65a99",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4096.000162495347"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jsq_fd = det(fd_jacobian(M -> M^2, M))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ce6c37e6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8.413470707960181e-6"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jsin_fd = det(fd_jacobian(sin, M))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d58d61d",
   "metadata": {},
   "source": [
    "So, we get **about 939.0, 4096, and -8.41×10⁻⁶**, respectively.\n",
    "\n",
    "We'll check these numbers against the analytical formula in part (ii)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47c10b6f",
   "metadata": {},
   "source": [
    "## (i) By automatic differentiation\n",
    "\n",
    "Alternatively, we can try to do this by AD:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "0738ca9b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import ForwardDiff, Zygote"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aec6524b",
   "metadata": {},
   "source": [
    "We'll start with the easy case of $M^2$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "39c58fd2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Int64}:\n",
       " 0  1  4  1  0  0  4  0  0\n",
       " 1  0  1  0  1  0  0  4  0\n",
       " 4  1  0  0  0  1  0  0  4\n",
       " 1  0  0  0  1  4  1  0  0\n",
       " 0  1  0  1  0  1  0  1  0\n",
       " 0  0  1  4  1  0  0  0  1\n",
       " 4  0  0  1  0  0  0  1  4\n",
       " 0  4  0  0  1  0  1  0  1\n",
       " 0  0  4  0  0  1  4  1  0"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jsq = ForwardDiff.jacobian(M -> M^2, M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "0c68139b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4096.0"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "det(Jsq)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2f195bf1",
   "metadata": {},
   "source": [
    "Unfortunately, `ForwardDiff` doesn't support the matrix-exponential function because the Julia linear-algebra library's `exp` function cannot handle matrices full of dual numbers.\n",
    "\n",
    "(See also the discussion in [ForwardDiff.jl#174](https://github.com/JuliaDiff/ForwardDiff.jl/issues/174).)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "8dcaae81",
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "MethodError: no method matching exp(::Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9}})\n\u001b[0mClosest candidates are:\n\u001b[0m  exp(\u001b[91m::Union{Float16, Float32, Float64}\u001b[39m) at special/exp.jl:326\n\u001b[0m  exp(\u001b[91m::StridedMatrix{var\"#s885\"} where var\"#s885\"<:Union{Float32, Float64, ComplexF32, ComplexF64}\u001b[39m) at /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:569\n\u001b[0m  exp(\u001b[91m::StridedMatrix{var\"#s885\"} where var\"#s885\"<:Union{Integer, Complex{<:Integer}}\u001b[39m) at /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:570\n\u001b[0m  ...",
     "output_type": "error",
     "traceback": [
      "MethodError: no method matching exp(::Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9}})\n\u001b[0mClosest candidates are:\n\u001b[0m  exp(\u001b[91m::Union{Float16, Float32, Float64}\u001b[39m) at special/exp.jl:326\n\u001b[0m  exp(\u001b[91m::StridedMatrix{var\"#s885\"} where var\"#s885\"<:Union{Float32, Float64, ComplexF32, ComplexF64}\u001b[39m) at /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:569\n\u001b[0m  exp(\u001b[91m::StridedMatrix{var\"#s885\"} where var\"#s885\"<:Union{Integer, Complex{<:Integer}}\u001b[39m) at /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:570\n\u001b[0m  ...",
      "",
      "Stacktrace:",
      " [1] vector_mode_dual_eval!(f::typeof(exp), cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9}}}, x::Matrix{Int64})",
      "   @ ForwardDiff ~/.julia/packages/ForwardDiff/QdStj/src/apiutils.jl:37",
      " [2] vector_mode_jacobian(f::typeof(exp), x::Matrix{Int64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9}}})",
      "   @ ForwardDiff ~/.julia/packages/ForwardDiff/QdStj/src/jacobian.jl:148",
      " [3] jacobian(f::Function, x::Matrix{Int64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9}}}, ::Val{true})",
      "   @ ForwardDiff ~/.julia/packages/ForwardDiff/QdStj/src/jacobian.jl:21",
      " [4] jacobian(f::Function, x::Matrix{Int64}, cfg::ForwardDiff.JacobianConfig{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9, Matrix{ForwardDiff.Dual{ForwardDiff.Tag{typeof(exp), Int64}, Int64, 9}}}) (repeats 2 times)",
      "   @ ForwardDiff ~/.julia/packages/ForwardDiff/QdStj/src/jacobian.jl:19",
      " [5] top-level scope",
      "   @ In[20]:1"
     ]
    }
   ],
   "source": [
    "ForwardDiff.jacobian(exp, M)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11b343c9",
   "metadata": {},
   "source": [
    "Let's try the [Zygote.jl package](https://github.com/FluxML/Zygote.jl)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "d9a3498e",
   "metadata": {},
   "outputs": [
    {
     "ename": "LoadError",
     "evalue": "Mutating arrays is not supported -- called setindex!(Matrix{Float64}, ...)\nThis error occurs when you ask Zygote to differentiate operations that change\nthe elements of arrays in place (e.g. setting values with x .= ...)\n\nPossible fixes:\n- avoid mutating operations (preferred)\n- or read the documentation and solutions for this error\n  https://fluxml.ai/Zygote.jl/latest/limitations\n",
     "output_type": "error",
     "traceback": [
      "Mutating arrays is not supported -- called setindex!(Matrix{Float64}, ...)\nThis error occurs when you ask Zygote to differentiate operations that change\nthe elements of arrays in place (e.g. setting values with x .= ...)\n\nPossible fixes:\n- avoid mutating operations (preferred)\n- or read the documentation and solutions for this error\n  https://fluxml.ai/Zygote.jl/latest/limitations\n",
      "",
      "Stacktrace:",
      "  [1] error(s::String)",
      "    @ Base ./error.jl:35",
      "  [2] _throw_mutation_error(f::Function, args::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/lib/array.jl:86",
      "  [3] (::Zygote.var\"#391#392\"{Matrix{Float64}})(#unused#::Nothing)",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/lib/array.jl:98",
      "  [4] (::Zygote.var\"#2488#back#393\"{Zygote.var\"#391#392\"{Matrix{Float64}}})(Δ::Nothing)",
      "    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67",
      "  [5] Pullback",
      "    @ /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:478 [inlined]",
      "  [6] Pullback",
      "    @ /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/matmul.jl:474 [inlined]",
      "  [7] (::typeof(∂(copytri!)))(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface2.jl:0",
      "  [8] Pullback",
      "    @ /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:628 [inlined]",
      "  [9] (::typeof(∂(exp!)))(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface2.jl:0",
      " [10] Pullback",
      "    @ /Applications/Julia-1.8.app/Contents/Resources/julia/share/julia/stdlib/v1.8/LinearAlgebra/src/dense.jl:570 [inlined]",
      " [11] (::typeof(∂(exp)))(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface2.jl:0",
      " [12] #208",
      "    @ ~/.julia/packages/Zygote/AS0Go/src/lib/lib.jl:206 [inlined]",
      " [13] (::Zygote.var\"#2066#back#210\"{Zygote.var\"#208#209\"{Tuple{Tuple{Nothing}}, typeof(∂(exp))}})(Δ::Matrix{Float64})",
      "    @ Zygote ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67",
      " [14] Pullback",
      "    @ ./operators.jl:1035 [inlined]",
      " [15] (::typeof(∂(#_#95)))(Δ::Vector{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface2.jl:0",
      " [16] (::Zygote.var\"#208#209\"{Tuple{Tuple{Nothing, Nothing}, Tuple{Nothing}}, typeof(∂(#_#95))})(Δ::Vector{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/lib/lib.jl:206",
      " [17] #2066#back",
      "    @ ~/.julia/packages/ZygoteRules/AIbCs/src/adjoint.jl:67 [inlined]",
      " [18] Pullback",
      "    @ ./operators.jl:1033 [inlined]",
      " [19] (::typeof(∂(ComposedFunction{typeof(Zygote._jvec), typeof(exp)}(Zygote._jvec, exp))))(Δ::Vector{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface2.jl:0",
      " [20] (::Zygote.var\"#60#61\"{typeof(∂(ComposedFunction{typeof(Zygote._jvec), typeof(exp)}(Zygote._jvec, exp)))})(Δ::Vector{Float64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/compiler/interface.jl:45",
      " [21] withjacobian(f::Function, args::Matrix{Int64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/lib/grad.jl:150",
      " [22] jacobian(f::Function, args::Matrix{Int64})",
      "    @ Zygote ~/.julia/packages/Zygote/AS0Go/src/lib/grad.jl:128",
      " [23] top-level scope",
      "    @ In[21]:1"
     ]
    }
   ],
   "source": [
    "Zygote.jacobian(exp, M)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "01e0c8e2",
   "metadata": {},
   "source": [
    "Nope!\n",
    "\n",
    "Maybe if we try to implement the matrix exponential ourselves using eigenvalues/vectors, which should be good for symmetric matrices at least:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "82bb775a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "fM (generic function with 1 method)"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function fM(f, M)\n",
    "    F = eigen(M)\n",
    "    return F.vectors * Diagonal(f.(F.values)) / F.vectors\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "72c78c4d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " 22.9849   0.0  0.0  14.9964   2.1476   0.0  36.7379  10.9334   13.7714\n",
       "  7.49818  0.0  0.0  13.5087   4.37449  0.0  12.9649  12.4211    5.46671\n",
       " 18.369    0.0  0.0  12.9649   2.1476   0.0  36.7562  12.9649   18.369\n",
       "  7.49818  0.0  0.0  13.5087   4.37449  0.0  12.9649  12.4211    5.46671\n",
       "  2.1476   0.0  0.0   8.74898  4.13663  0.0   4.2952   8.74898   2.1476\n",
       "  5.46671  0.0  0.0  12.4211   4.37449  0.0  12.9649  13.5087    7.49818\n",
       " 18.369    0.0  0.0  12.9649   2.1476   0.0  36.7562  12.9649   18.369\n",
       "  5.46671  0.0  0.0  12.4211   4.37449  0.0  12.9649  13.5087    7.49818\n",
       " 13.7714   0.0  0.0  10.9334   2.1476   0.0  36.7379  14.9964   22.9849"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Zygote.jacobian(M -> fM(exp, M), M)[1]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5255eed3",
   "metadata": {},
   "source": [
    "Success!  Hmm, it doesn't look right though — those columns of zeros look wrong.\n",
    "\n",
    "It turns out that this is a [bug in Zygote](https://github.com/FluxML/Zygote.jl/issues/1369) for computing derivatives of eigenvalues when the matrix happen to be real-symmetric/Hermitian.\n",
    "\n",
    "We can work around the problem by slightly perturbing M to make it nonsymmetric:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "182c9233",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " 22.9849    7.49818  18.369     7.49818  …  18.369     5.46671  13.7714\n",
       "  7.49818  11.3611    7.49818   2.1476       5.46671  10.2735    5.46671\n",
       " 18.369     7.49818  22.9849    5.46671     13.7714    5.46671  18.369\n",
       "  7.49818   2.1476    5.46671  11.3611       7.49818   2.1476    5.46671\n",
       "  2.1476    4.37449   2.1476    4.37449      2.1476    4.37449   2.1476\n",
       "  5.46671   2.1476    7.49818  10.2735   …   5.46671   2.1476    7.49818\n",
       " 18.369     5.46671  13.7714    7.49818     22.9849    7.49818  18.369\n",
       "  5.46671  10.2735    5.46671   2.1476       7.49818  11.3611    7.49818\n",
       " 13.7714    5.46671  18.369     5.46671     18.369     7.49818  22.9849"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jexp = Zygote.jacobian(M -> fM(exp, M), M + [0 0 1e-13; 0 0 0; 0 0 0])[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "8687f2d8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "939.0589449936064"
      ]
     },
     "execution_count": 25,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "det(Jexp)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "3b715a07",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Int64}:\n",
       " 0  1  4  1  0  0  4  0  0\n",
       " 1  0  1  0  1  0  0  4  0\n",
       " 4  1  0  0  0  1  0  0  4\n",
       " 1  0  0  0  1  4  1  0  0\n",
       " 0  1  0  1  0  1  0  1  0\n",
       " 0  0  1  4  1  0  0  0  1\n",
       " 4  0  0  1  0  0  0  1  4\n",
       " 0  4  0  0  1  0  1  0  1\n",
       " 0  0  4  0  0  1  4  1  0"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jsq = Zygote.jacobian(M -> M^2, M)[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "e920fcc3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4096.0"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "det(Jsq)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "920b0c12",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " -0.327595   -0.0100503   0.107194   …   0.107194   -0.0369293  -0.111661\n",
       " -0.0100503  -0.180199   -0.0100503     -0.0369293   0.14325    -0.0369293\n",
       "  0.107194   -0.0100503  -0.327595      -0.111661   -0.0369293   0.107194\n",
       " -0.0100503   0.0357343  -0.0369293     -0.0100503   0.0357343  -0.0369293\n",
       "  0.0357343  -0.189917    0.0357343      0.0357343  -0.189917    0.0357343\n",
       " -0.0369293   0.0357343  -0.0100503  …  -0.0369293   0.0357343  -0.0100503\n",
       "  0.107194   -0.0369293  -0.111661      -0.327595   -0.0100503   0.107194\n",
       " -0.0369293   0.14325    -0.0369293     -0.0100503  -0.180199   -0.0100503\n",
       " -0.111661   -0.0369293   0.107194       0.107194   -0.0100503  -0.327595"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jsin = Zygote.jacobian(M -> fM(sin, M), M + [0 0 1e-13; 0 0 0; 0 0 0])[1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "aef5c10a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8.413463127258828e-6"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "det(Jsin)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d22e3794",
   "metadata": {},
   "source": [
    "So, we eventually got determinants of **about 939.1, exactly 4096, and about -8.41×10⁻⁶**, respectively, pretty close to our finite-difference answers!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb5bd31c",
   "metadata": {},
   "source": [
    "## (ii) Analytical formula\n",
    "\n",
    "Now we'll compare to the analytical formula:\n",
    "$$\n",
    "\\frac{\\prod_{i < j} |f(\\lambda_i)-\n",
    "f(\\lambda_j)|^2}{\\prod_{i < j} |\\lambda_i-\\lambda_j|^2}\n",
    "\\prod_i f'(\\lambda_i)\n",
    "$$\n",
    "or equivalently\n",
    "$$\n",
    "\\prod_{\\lambda_i \\ne \\lambda_j} \\frac{|f(\\lambda_i)-\n",
    "f(\\lambda_j)|}{|\\lambda_i-\\lambda_j|}\n",
    "\\prod_i f'(\\lambda_i)\n",
    "$$\n",
    "\n",
    "with a little help from the [`product` iterator](https://docs.julialang.org/en/v1/base/iterators/#Base.Iterators.product):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "75ed8d35",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Jdet_analytic (generic function with 1 method)"
      ]
     },
     "execution_count": 30,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function Jdet_analytic(f, M)\n",
    "    λ = eigvals(M)\n",
    "    return prod(λ -> ForwardDiff.derivative(f, λ), λ) *\n",
    "           prod(Iterators.product(λ, λ)) do ((λᵢ,λⱼ))\n",
    "              λᵢ ≠ λⱼ ? abs((f(λᵢ)-f(λⱼ)) / (λᵢ-λⱼ)) : 1.0\n",
    "           end\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "2707a3c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "939.058944993412"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jexp_analytic = Jdet_analytic(exp, M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "id": "7bc84a93",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "4095.999999999958"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jsq_analytic = Jdet_analytic(M -> M^2, M)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "2372177f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "8.413463127260002e-6"
      ]
     },
     "execution_count": 33,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Jsin_analytic = Jdet_analytic(sin, M)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "341493f5",
   "metadata": {},
   "source": [
    "That is, we got determinants of **about 939.1, 4096, and -8.41×10⁻⁶**, respectively, pretty close to our answers above!\n",
    "\n",
    "Let's compare them more quantitatively by computing the relative errors in the finite-difference approximations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "id": "a98b99db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2.8896516129601107e-5, 3.9671725682311325e-8, 9.010201939337017e-7)"
      ]
     },
     "execution_count": 34,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(Jexp_fd, Jexp_analytic),\n",
    "relerr(Jsq_fd, Jsq_analytic),\n",
    "relerr(Jsin_fd, Jsin_analytic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8d49cd85",
   "metadata": {},
   "source": [
    "So, the finite-difference results were accurate to about 5 digits, 8 digits, and 6 digits, respectively.  Not too bad!\n",
    "\n",
    "And now computing the relative errors in the Zygote AD results:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "id": "cabe756c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(2.0702054278959802e-13, 1.0214051826551544e-14, 1.3953679324874998e-13)"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(det(Jexp), Jexp_analytic),\n",
    "relerr(det(Jsq), Jsq_analytic),\n",
    "relerr(det(Jsin), Jsin_analytic)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f27270a",
   "metadata": {},
   "source": [
    "The Zygote results were accurate to about 13 digits, which is pretty good considering that we \"fudged\" $M$ by $10^{-13}$ to work around the Zygote bug for symmetric $M$."
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.8.0",
   "language": "julia",
   "name": "julia-1.8"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}