{
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "# Sensitivity Analysis\n",
    "\n",
    "In order to facilitate gradient-based optimization, all linear solves in GXBeam have been\n",
    "overloaded to support automatic differentiation using the ImplicitAD package.  Automatic\n",
    "differentiation using ForwardDiff and ReverseDiff should therefore work with without any\n",
    "major headaches, as long as you initialize the internal storage with the correct type.\n",
    "(e.g. `system = DynamicSystem(TF, assembly)` where `TF` is an appropriate floating point\n",
    "type).\n",
    "\n",
    "For example, consider the Cantilever with a Tip Moment example.  Suppose\n",
    "we were interested in the sensitivity of tip x and y-displacement with respect to\n",
    "the nondimensional tip moment $\\lambda$ when $\\lambda=1$.  These sensitivites may\n",
    "be computed as follows:"
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  116.064 ms (52133 allocations: 5.15 MiB)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "2×1 Matrix{Float64}:\n -12.025539968259286\n  -7.659921842013467"
     },
     "metadata": {},
     "execution_count": 1
    }
   ],
   "cell_type": "code",
   "source": [
    "using GXBeam, LinearAlgebra\n",
    "import ForwardDiff # for forward-mode automatic differentiation\n",
    "import ReverseDiff # for reverse-mode automatic differentiation\n",
    "using BenchmarkTools # for benchmarking function performance\n",
    "\n",
    "L = 12 # inches\n",
    "h = w = 1 # inches\n",
    "E = 30e6 # lb/in^4 Young's Modulus\n",
    "\n",
    "A = h*w\n",
    "Iyy = w*h^3/12\n",
    "Izz = w^3*h/12\n",
    "\n",
    "# create points\n",
    "nelem = 16\n",
    "x = range(0, L, length=nelem+1)\n",
    "y = zero(x)\n",
    "z = zero(x)\n",
    "points = [[x[i],y[i],z[i]] for i = 1:length(x)]\n",
    "\n",
    "# index of endpoints of each beam element\n",
    "start = 1:nelem\n",
    "stop = 2:nelem+1\n",
    "\n",
    "# compliance matrix for each beam element\n",
    "compliance = fill(Diagonal([1/(E*A), 0, 0, 0, 1/(E*Iyy), 1/(E*Izz)]), nelem)\n",
    "\n",
    "# create assembly of interconnected nonlinear beams\n",
    "assembly = Assembly(points, start, stop, compliance=compliance)\n",
    "\n",
    "# construct objective function\n",
    "objfun1 = (p) -> begin\n",
    "\n",
    "    # non-dimensional tip moment\n",
    "    λ = p[1]\n",
    "\n",
    "    # dimensionalized tip moment\n",
    "    m = pi*E*Iyy/L\n",
    "    M = λ*m\n",
    "\n",
    "    # prescribed conditions\n",
    "    prescribed_conditions = Dict(\n",
    "        # fixed left side\n",
    "        1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0, theta_z=0),\n",
    "        # moment on right side\n",
    "        nelem+1 => PrescribedConditions(Mz = M)\n",
    "    )\n",
    "\n",
    "    # initialize internal storage with the correct type\n",
    "    system = StaticSystem(eltype(p), assembly)\n",
    "\n",
    "    # perform static analysis\n",
    "    system, state, converged = static_analysis!(system, assembly; prescribed_conditions)\n",
    "\n",
    "    # return the desired outputs\n",
    "    return [state.points[end].u[1], state.points[end].u[2]]\n",
    "end\n",
    "\n",
    "# compute sensitivities using ForwardDiff with λ = 1.0\n",
    "@btime ForwardDiff.jacobian(objfun1, [1.0])"
   ],
   "metadata": {},
   "execution_count": 1
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  134.525 ms (2201447 allocations: 96.70 MiB)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "2×1 Matrix{Float64}:\n -12.025539968259285\n  -7.659921842013468"
     },
     "metadata": {},
     "execution_count": 2
    }
   ],
   "cell_type": "code",
   "source": [
    "# compute sensitivities using ReverseDiff with λ = 1.0\n",
    "@btime ReverseDiff.jacobian(objfun1, [1.0])"
   ],
   "metadata": {},
   "execution_count": 2
  },
  {
   "cell_type": "markdown",
   "source": [
    "Advanced users, however, may wish to use overloaded versions of each nonlinear solve\n",
    "in order to further decrease the total computational costs associated with obtaining\n",
    "design sensitivites.  Overloading the nonlinear solver also significantly reduces the\n",
    "memory requirements associated with using ReverseDiff. Using these overloads, however,\n",
    "requires that the user provide the parameter function `parameters = pfunc(p, t)` and\n",
    "associated parameter vector `p`.  As described in the documentation for each analysis\n",
    "type, the `pfunc` function returns a named tuple which contains updated arguments for\n",
    "the analysis, based on the contents of the parameter vector `p` and the current time `t`."
   ],
   "metadata": {}
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  41.044 ms (13729 allocations: 3.13 MiB)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "2×1 Matrix{Float64}:\n -12.025539968209076\n  -7.659921841992032"
     },
     "metadata": {},
     "execution_count": 3
    }
   ],
   "cell_type": "code",
   "source": [
    "# construct pfunc to overwrite prescribed conditions\n",
    "pfunc = (p, t) -> begin\n",
    "\n",
    "    # non-dimensional tip moment\n",
    "    λ = p[1]\n",
    "\n",
    "    # dimensionalized tip moment\n",
    "    m = pi*E*Iyy/L\n",
    "    M = λ*m\n",
    "\n",
    "    # create dictionary of prescribed conditions\n",
    "    prescribed_conditions = Dict(\n",
    "        # fixed left side\n",
    "        1 => PrescribedConditions(ux=0, uy=0, uz=0, theta_x=0, theta_y=0, theta_z=0),\n",
    "        # moment on right side\n",
    "        nelem+1 => PrescribedConditions(Mz = M)\n",
    "    )\n",
    "\n",
    "    # return named tuple with new arguments\n",
    "    return (; prescribed_conditions=prescribed_conditions)\n",
    "end\n",
    "\n",
    "# construct objective function\n",
    "objfun2 = (p) -> begin\n",
    "\n",
    "    # perform static analysis\n",
    "    system, state, converged = static_analysis(assembly; pfunc, p)\n",
    "\n",
    "    # return the desired outputs\n",
    "    return [state.points[end].u[1], state.points[end].u[2]]\n",
    "end\n",
    "\n",
    "# compute sensitivities using ForwardDiff with λ = 1.0\n",
    "@btime ForwardDiff.jacobian(objfun2, [1.0])"
   ],
   "metadata": {},
   "execution_count": 3
  },
  {
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "  6.290 ms (13863 allocations: 3.66 MiB)\n"
     ]
    },
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": "2×1 Matrix{Float64}:\n -12.025539968209076\n  -7.659921841992033"
     },
     "metadata": {},
     "execution_count": 4
    }
   ],
   "cell_type": "code",
   "source": [
    "# compute sensitivities using ReverseDiff with λ = 1.0\n",
    "@btime ReverseDiff.jacobian(objfun2, [1.0])"
   ],
   "metadata": {},
   "execution_count": 4
  },
  {
   "cell_type": "markdown",
   "source": [
    "---\n",
    "\n",
    "*This notebook was generated using [Literate.jl](https://github.com/fredrikekre/Literate.jl).*"
   ],
   "metadata": {}
  }
 ],
 "nbformat_minor": 3,
 "metadata": {
  "language_info": {
   "file_extension": ".jl",
   "mimetype": "application/julia",
   "name": "julia",
   "version": "1.9.1"
  },
  "kernelspec": {
   "name": "julia-1.9",
   "display_name": "Julia 1.9.1",
   "language": "julia"
  }
 },
 "nbformat": 4
}