hw2sol.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "eb604661",
   "metadata": {},
   "source": [
    "Similar to homework 1, let's define a relative-error function for quantifying the error from finite differences:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "48cd6341",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "relerr (generic function with 1 method)"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "using LinearAlgebra # for norm, kron, I(n) etc.\n",
    "\n",
    "relerr(approx, exact) = norm(approx - exact) / norm(exact)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3ccad3f3",
   "metadata": {},
   "source": [
    "Let's also define the $\\otimes$ operator (type it by `\\otimes` followed by TAB) to be the Kronecker product (the `kron` function in the `LinearAlgebra` library) so we can use math notation `A ⊗ B` instead of `kron(A, B)`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "24b5fd38",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "kron (generic function with 21 methods)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "const ⊗ = kron"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43df9092",
   "metadata": {},
   "source": [
    "We'll also load the Zygote AD library for analytical comparisons:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0a1fc52b",
   "metadata": {},
   "outputs": [],
   "source": [
    "using Zygote"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c8e37b5",
   "metadata": {},
   "source": [
    "# Problem 2\n",
    "\n",
    "Here, $f(A) = \\sqrt{A}$, and the homework says the Jacobian (acting on $\\operatorname{vec}(A)$) should be $(I\\otimes \\sqrt{A} + \\sqrt{A}^T \\otimes I)^{-1}$.  Let's try it out for a random positive-definite $A$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "63dc982c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×1 Matrix{Float64}:\n",
       " 1.3731197861011764"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n = 1 # should be big enough to be interesting\n",
    "B = randn(n,n)\n",
    "A = B'B # random positive definite"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "993f6868",
   "metadata": {},
   "source": [
    "The matrix square root is just `sqrt(A)` in Julia:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cb14ff97",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# check that sqrt(A)² ≈ A, up to roundoff errors:\n",
    "relerr(sqrt(A)^2, A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "0b3cb733",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1×1 Matrix{Float64}:\n",
       " 0.42669326873498054"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Jacobian:\n",
    "J = (I(n) ⊗ sqrt(A) + sqrt(A)' ⊗ I(n))^-1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3b89a4f5",
   "metadata": {},
   "source": [
    "Now, let's check this against finite difference for a small random $dA$:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "0d188b50",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3.726301788530775e-8"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dA = randn(n,n) * 1e-8\n",
    "\n",
    "relerr( vec(sqrt(A+dA) - sqrt(A)),  # finite-difference directional derivative\n",
    "        J * vec(dA) )               # vs. exact expression"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5574a08",
   "metadata": {},
   "source": [
    "Hooray, it worked!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cd27f68",
   "metadata": {},
   "source": [
    "# Problem 3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3d8f59a",
   "metadata": {},
   "source": [
    "Let's define our function $f(p)$ as in the problem.  We'll pass the extra parameters $A_0$ etcetera explicitly (or we could use global variables):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "c2b586ab",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "f (generic function with 1 method)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function f(p, A₀, a, B₀, b, F)\n",
    "    A = A₀ + diagm(p)\n",
    "    x = A \\ a  # A⁻¹a\n",
    "    B = B₀ + diagm(x.*x)\n",
    "    y = B \\ b  # B⁻¹b\n",
    "    return y'*F*y\n",
    "end"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1ba4a2c9",
   "metadata": {},
   "source": [
    "Now, we'll pick some random parameters and try it out:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "2147c2b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "49.975441720918774"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "n = 5\n",
    "A₀ = randn(n,n)\n",
    "a = randn(5)\n",
    "B₀ = randn(n,n)\n",
    "b = randn(5)\n",
    "F = randn(n,n); F = F + F'\n",
    "p = rand(5)\n",
    "f(p, A₀, a, B₀, b, F)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2484c042",
   "metadata": {},
   "source": [
    "Now, let's implement the *manual* \"adjoint\" gradient from the solutions, using the same notation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "002b28b3",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "∇f (generic function with 1 method)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "function ∇f(p, A₀, a, B₀, b, F)\n",
    "    # the forward solution, copied from above\n",
    "    #    (note that in serious computation we would want to re-use this from f)\n",
    "    A = A₀ + diagm(p)\n",
    "    x = A \\ a  # A⁻¹a\n",
    "    B = B₀ + diagm(x.*x)\n",
    "    y = B \\ b  # B⁻¹b\n",
    "    \n",
    "    # the reverse-mode gradient\n",
    "    g′ᵀ = 2F*y         # step (i)\n",
    "    u = B' \\ -g′ᵀ      # step (ii): adjoint problem 1\n",
    "    w = 2u .* x .* y   # step (iii)\n",
    "    z = A' \\ -w        # step (iv): adjoint problem 2\n",
    "    return z .* x      # step (v): ∇f\n",
    "end"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5fa14097",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5-element Vector{Float64}:\n",
       "  262.55155103380923\n",
       " -273.8595112733488\n",
       "   96.32881037878508\n",
       "  217.43756038923235\n",
       " -171.79342086396048"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "∇f(p, A₀, a, B₀, b, F)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c1b030e",
   "metadata": {},
   "source": [
    "## problem 3b\n",
    "\n",
    "First, we'll compare it against finite differences for a random small `dp`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "657849bd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.428248665134681e-7"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dp = randn(n) * 1e-8\n",
    "\n",
    "relerr( f(p+dp, A₀, a, B₀, b, F) - f(p, A₀, a, B₀, b, F),  # finite difference\n",
    "        ∇f(p, A₀, a, B₀, b, F)'dp) # exact directional derivative"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b5371853",
   "metadata": {},
   "source": [
    "Hooray, it matches!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cae7ae8e",
   "metadata": {},
   "source": [
    "## problem 3c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f15be256",
   "metadata": {},
   "source": [
    "To compute $\\nabla f$ with `Zygote`, we have to give Zygote a function of a *single* parameter vector `p` that we want to differentiate with respect to, along with the point `p` at which we want the derivative.  To do that, we will define an [anonymous function](https://docs.julialang.org/en/v1/manual/functions/#man-anonymous-functions) with `p -> ...` that captures the other parameters (also called a [closure](https://en.wikipedia.org/wiki/Closure_(computer_programming)) in computer science):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "2329ed32",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "([262.5515510338092, -273.8595112733488, 96.32881037878506, 217.43756038923235, -171.7934208639604],)"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Zygote.gradient(p -> f(p, A₀, a, B₀, b, F), p)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "09a8a25a",
   "metadata": {},
   "source": [
    "(Zygote returns a 1-component tuple of outputs, because it can potentially differentiate with respect to multiple arguments, though here we are just asking for 1.)\n",
    "\n",
    "The above looks pretty good if we \"eyeball\" it compared to `∇f` above.  Let's compare it quantitatively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c0c3d50f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "2.1572017312010847e-16"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(∇f(p, A₀, a, B₀, b, F),                              # manual ∇f\n",
    "       Zygote.gradient(p -> f(p, A₀, a, B₀, b, F), p)[1])   #  vs AD"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33c8c7d5",
   "metadata": {},
   "source": [
    "Hooray, it matches up to the limits of roundoff error (to essentially [machine precision](https://en.wikipedia.org/wiki/Machine_epsilon))!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cb40856",
   "metadata": {},
   "source": [
    "# Problem 4c"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8834f86c",
   "metadata": {},
   "source": [
    "Demonstrate numerically that  $d(e^A) = \\sum_{k=0}^\\infty  \\! \\! \\frac{1}{k!} (\\sum_{\\ell=0}^{k-1} (A^T)^{k-\\ell-1} \\otimes A^\\ell )dA$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "36ca240e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3×3 Matrix{Float64}:\n",
       " 0.989401  0.372355  0.475427\n",
       " 0.94298   0.68495   0.646221\n",
       " 0.794569  0.13317   0.353663"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "A = rand(3,3)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "df6620d8",
   "metadata": {},
   "source": [
    "### Using ForwardDiff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "a7dfb9f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "using ForwardDiff"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "d3176e9d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "e (generic function with 1 method)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "e(A) = sum(A^k/factorial(k) for k=0:20) # hmm exp doesn't work, i'll go to k=20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c1536121",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.2971024122201743e-15"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(e(A), exp(A))  # check that our sum matches the built-in exp(A) function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "8b7817cb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " 3.35707   0.5389     0.688955   1.50121   …  1.04456    0.123629   0.158014\n",
       " 1.50121   2.82372    0.907286   0.466706     0.338304   0.92634    0.209389\n",
       " 1.04456   0.265598   2.45771    0.338304     0.244335   0.0576795  0.83731\n",
       " 0.5389    0.0624931  0.0798733  2.82372      0.265598   0.0290732  0.0371575\n",
       " 0.170853  0.479248   0.105876   1.3384       0.079272   0.237991   0.0492994\n",
       " 0.123629  0.0290732  0.434115   0.92634   …  0.0576795  0.0134119  0.216817\n",
       " 0.688955  0.0798733  0.102087   0.907286     2.45771    0.434115   0.555028\n",
       " 0.218367  0.612715   0.135322   0.289541     1.21467    2.02443    0.729772\n",
       " 0.158014  0.0371575  0.555028   0.209389     0.83731    0.216817   1.73404"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "J_AD = ForwardDiff.jacobian(e,A)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "84969d50",
   "metadata": {},
   "source": [
    "### .. or using Finite Differences"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "7b35a6b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " 3.35707   0.538899   0.688955   1.50121   …  1.04456    0.123629   0.158014\n",
       " 1.50121   2.82372    0.907286   0.466706     0.338304   0.92634    0.209389\n",
       " 1.04456   0.265598   2.45771    0.338304     0.244335   0.0576795  0.83731\n",
       " 0.5389    0.062493   0.0798732  2.82372      0.265598   0.0290732  0.0371576\n",
       " 0.170853  0.479248   0.105875   1.3384       0.0792719  0.237991   0.0492994\n",
       " 0.123629  0.0290732  0.434115   0.92634   …  0.0576795  0.0134118  0.216817\n",
       " 0.688955  0.0798732  0.102087   0.907286     2.45771    0.434115   0.555028\n",
       " 0.218367  0.612715   0.135321   0.289541     1.21467    2.02443    0.729772\n",
       " 0.158014  0.0371575  0.555028   0.209389     0.83731    0.216817   1.73404"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ϵ = 1e-8\n",
    "J = zeros(9,0)  # initialize 9x9 Jacobian with 9 rows and no columns\n",
    "for j=1:3, i=1:3\n",
    "    dA = zeros(3,3)\n",
    "    dA[i,j] = ϵ  # perturb the (i,j) entry only\n",
    "    df = exp(A+dA)-exp(A) # see the perturbed exp\n",
    "    J = [J vec(df)] # append this to J\n",
    "end\n",
    "J_FD = J/ϵ"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6973563",
   "metadata": {},
   "source": [
    "### The theoretical answer summing to 20"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "aabb597f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " 3.35707   0.5389     0.688955   1.50121   …  1.04456    0.123629   0.158014\n",
       " 1.50121   2.82372    0.907286   0.466706     0.338304   0.92634    0.209389\n",
       " 1.04456   0.265598   2.45771    0.338304     0.244335   0.0576795  0.83731\n",
       " 0.5389    0.0624931  0.0798733  2.82372      0.265598   0.0290732  0.0371575\n",
       " 0.170853  0.479248   0.105876   1.3384       0.079272   0.237991   0.0492994\n",
       " 0.123629  0.0290732  0.434115   0.92634   …  0.0576795  0.0134119  0.216817\n",
       " 0.688955  0.0798733  0.102087   0.907286     2.45771    0.434115   0.555028\n",
       " 0.218367  0.612715   0.135322   0.289541     1.21467    2.02443    0.729772\n",
       " 0.158014  0.0371575  0.555028   0.209389     0.83731    0.216817   1.73404"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# written as nested sum calls:\n",
    "sum(sum( (A')^(k-ℓ-1) ⊗ A^ℓ for ℓ=0:(k-1))/factorial(k) for k=1:20)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "75eb21e5",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "9×9 Matrix{Float64}:\n",
       " 3.35707   0.5389     0.688955   1.50121   …  1.04456    0.123629   0.158014\n",
       " 1.50121   2.82372    0.907286   0.466706     0.338304   0.92634    0.209389\n",
       " 1.04456   0.265598   2.45771    0.338304     0.244335   0.0576795  0.83731\n",
       " 0.5389    0.0624931  0.0798733  2.82372      0.265598   0.0290732  0.0371575\n",
       " 0.170853  0.479248   0.105876   1.3384       0.079272   0.237991   0.0492994\n",
       " 0.123629  0.0290732  0.434115   0.92634   …  0.0576795  0.0134119  0.216817\n",
       " 0.688955  0.0798733  0.102087   0.907286     2.45771    0.434115   0.555028\n",
       " 0.218367  0.612715   0.135322   0.289541     1.21467    2.02443    0.729772\n",
       " 0.158014  0.0371575  0.555028   0.209389     0.83731    0.216817   1.73404"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# same thing written another way:\n",
    "J_20 = sum( (A')^(k-ℓ-1) ⊗ A^ℓ / factorial(k) for ℓ=0:20, k=1:20 if k>ℓ)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bd39711b",
   "metadata": {},
   "source": [
    "Looks good let's check the match quantitatively:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "d3c6f007",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.3738128345184809e-15"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(J_20, J_AD)  # should match to nearly machine precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6c1d30c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.293987776257493e-7"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "relerr(J_20, J_FD)  # should match to ≈ 7 digits"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49ebb275",
   "metadata": {},
   "source": [
    "Hooray, math works!"
   ]
  }
 ],
 "metadata": {
  "@webio": {
   "lastCommId": null,
   "lastKernelId": null
  },
  "kernelspec": {
   "display_name": "Julia 1.7.1",
   "language": "julia",
   "name": "julia-1.7"
  },
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.7.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}