AIFT_best_fit_example.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Y0MpmK1GSuCB"
      },
      "source": [
        "# Automatic Implicit Function Theorem\n",
        "## Least-squares example\n",
        "\n",
        "Here we augment the example in `AIFT_exact_fit_example.ipynb` with an introduction of a best-fit least-squares calibration instead of the exactly-solvable calibration problem in the main sheet\n",
        "\n",
        "We use the method fo Section 5.2 of the paper that produces approximate adjoints that are close to the actual ones as long as the least-squares problem is \"close\" to the exact fit, i.e. the regularization weight is small\n",
        "\n",
        "Of particular note in this example is how few change to the AIFT code are required to support the least-squares best-fit. Most of the extra code added is here because we apply regularization to the derivative of the fitting function which requires some more claculations, and not from the AIFT method"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lHcKUzrESuvd"
      },
      "outputs": [],
      "source": [
        "import matplotlib.pyplot as plt\n",
        "from scipy.optimize import least_squares\n",
        "\n",
        "from autograd import grad, jacobian # pip install autograd\n",
        "import autograd.numpy as np   # Thinly-wrapped version of Numpy to fascilitate autograd calcs\n",
        "\n",
        "# This is the regularization weight that is used later in the objective function to turn exact-fit into best-fit\n",
        "regularization_weight = 0.001"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "N08P2Hcc6VD3"
      },
      "source": [
        "## Define the polynomial value function and the interpolator function"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ss0Feg7Eyur"
      },
      "source": [
        "\n",
        "In this section we set up the problem that we will apply AAD to.\n",
        "\n",
        "The function **spolyval(...)** gives the values of the 'stretched polynomial' at times **ts[n]** given the coefficients ***coefs*** and weights $w$.\n",
        "\n",
        "The 'stretched polynomial' function is a (somewhat contrived but simple) example of an interpolator that fits all the points but which interpolation scheme, and the shape of the function between knots, depends on the weights $w$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YxT7YPWVFAK-"
      },
      "outputs": [],
      "source": [
        "def spolyval(coefs, ts, w):\n",
        "    '''\n",
        "    A 'stretched polynomial' function, a polynomial in wts,\n",
        "    where wts = w[0]*ts + w[1]*ts**2.\n",
        "    Wegiths w here control the shape of the function between knots.\n",
        "\n",
        "    coefs:  polynomial coefs\n",
        "    ts:     points where the function is evaluated\n",
        "    w:      weights to transform ts into wts\n",
        "    '''\n",
        "    tsw = w[0]*ts + w[1]*ts**2\n",
        "    val = 0.0\n",
        "    for n in range(len(coefs)):\n",
        "        val =  val + coefs[n] * tsw**n\n",
        "    return val\n",
        "\n",
        "def spolyval_derivative(coefs, ts, w):\n",
        "    # Derivative of spolyval()\n",
        "    tsw = w[0]*ts + w[1]*ts**2\n",
        "    tsw_der = w[0] + w[1]*ts*2\n",
        "\n",
        "    val = 0.0\n",
        "    for n in range(len(coefs)):\n",
        "        if (n>0):\n",
        "          val =  val + coefs[n] * tsw**(n-1) * n * tsw_der\n",
        "    return val"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lCiZxi5jFFdX"
      },
      "source": [
        "We simulate a typical programming pattern where auxilliary variables such as $w$\n",
        "come wrapped in various helpers, etc. This is not strictly necessary in this code but will be used later to illustrate important points. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eSiFOc9v6VD5"
      },
      "outputs": [],
      "source": [
        "\n",
        "\n",
        "class PricingHelper:\n",
        "     def __init__(self, w):\n",
        "         self.w_ = w\n",
        "         self.updatable = False\n",
        "         # If w is none we link w's to xs's in a particular way \n",
        "         # to introduce the extra dependence of the result of spoly_interp \n",
        "         # on xs via w (admittedly, somewhat artificially). The actual update \n",
        "         # happens in the update(...) function that the clients are supposed \n",
        "         # to call when the xs are known.\n",
        "         if w is None:\n",
        "            self.updatable = True\n",
        "\n",
        "     def update(self, xs, ts):\n",
        "        '''\n",
        "        Update the weights depending on the inputs ts (not used \n",
        "        in this example) and xs.\n",
        "        '''\n",
        "        if self.updatable:\n",
        "            self.w_ = np.array([1.0, np.sum(xs**2)])\n",
        "     \n",
        "     def spolyval(self,c,ts):\n",
        "       return spolyval(c, ts, self.w_) \n",
        "\n",
        "     def spolyval_derivative(self, c, ts):\n",
        "       return spolyval_derivative(c, ts, self.w_) * regularization_weight \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "t9k59d9VAYUV"
      },
      "source": [
        "Function **spoly_interp(...)** calculates the coefs by fitting spolyval to **ts,xs** and returns the value of **spolyval** at some other point **t**.\n",
        "Note how **w** is never seen inside the body of the function, all wrapped in **PricingHelper**.\n",
        "\n",
        "Executing **spoly_interp(...)** corresponds to computing the implicit function $C(x,w)$ from the paper.\n",
        "\n",
        "$\\Omega(c,x,w)$ in the paper is the concatenation of **spolyval(...)** and **spolyval_derivative(...)** multiplied by a regularization weight."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xTAy5EzKAYth"
      },
      "outputs": [],
      "source": [
        "    def spoly_interp(xs, ts, t,  pricing_helper):\n",
        "        '''\n",
        "        Best-fit a stretched polynomial to (ts,xs) and evaluate it at t\n",
        "        Here pricing_helper (via pricing_helper.w_) is defining \n",
        "        the interpolation between the knots. \n",
        "        '''\n",
        "        pricing_helper.update(xs,ts)\n",
        "    \n",
        "        # Note the change to obj_f that now applies a regularization weight to the derivatives of the fitting function at the t nodes\n",
        "        def obj_f(c, pricing_helper = pricing_helper):\n",
        "            return ( np.concatenate( (pricing_helper.spolyval(c, ts) - xs, pricing_helper.spolyval_derivative(c,ts) * regularization_weight) ))\n",
        "\n",
        "        x0 = np.zeros_like(ts)\n",
        "        res = least_squares(obj_f, x0)\n",
        "        return pricing_helper.spolyval(res.x, t)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MLoFFnQvFbiK"
      },
      "source": [
        "An example of applying **spoly_interp(...)**."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Sm7VufN96VD7"
      },
      "outputs": [],
      "source": [
        "# points we interpolate\n",
        "ts = np.array([0.,1,2,3,4])\n",
        "xs = np.array([2.,1,3,4,0])\n",
        "npts = len(xs)\n",
        "\n",
        "# the point at which we evaluate our interpolator\n",
        "t = 3.5 \n",
        "\n",
        "# We can try different values of w. 'None' is the default that triggers\n",
        "# the calculation w = w(x)\n",
        "# \n",
        "# w(xs) for the particular xs above is equal to [1.0,30.0] so \n",
        "# we can pass them directly and it will not affect the output value of \n",
        "# spoly_interp(...) but of course will affect the gradients\n",
        "\n",
        "# Uncomment one of these\n",
        "w_to_use = None\n",
        "#w_to_use = [1.0,30.0]\n",
        "\n",
        "# Set up the pricer helper\n",
        "pricing_helper = PricingHelper(w_to_use)\n",
        "\n",
        "# calculate the interpolated value\n",
        "v = spoly_interp(xs,ts,t, pricing_helper)\n",
        "print(v)\n",
        "\n",
        "# plot a graph to see what the interpolation function looks like\n",
        "t_fine = ts[0] + np.arange(101)/100*(ts[-1] - ts[0])\n",
        "v_fine = spoly_interp(xs,ts,t_fine, pricing_helper)\n",
        "plt.plot(t_fine, v_fine, '-', label = 'fitted interpolator')\n",
        "plt.plot(ts,xs,'o', label = 'interpolated points')\n",
        "plt.plot(t,v,'o',label = 'evaluation point')\n",
        "plt.legend(loc = 'best')\n",
        "\n",
        "plt.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qf8K-UDyBLj6"
      },
      "source": [
        "Now calculate the gradients using bumping."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "_ky9A-LPK4he"
      },
      "outputs": [],
      "source": [
        "eps = 1e-5\n",
        "grad_bump = np.zeros_like(xs)\n",
        "for n in range(len(xs)):\n",
        "    x1 = xs.copy()\n",
        "    x1[n] += eps\n",
        "    grad_bump[n] = (spoly_interp(x1, ts, t, pricing_helper) - spoly_interp(xs, ts, t, pricing_helper))/eps\n",
        "np.set_printoptions(precision=3)\n",
        "print(f'gradients by bumping = {grad_bump}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pFxZggYV6VD9"
      },
      "source": [
        "## Try `autograd` on poly_interp, 'differentiating' through the best-fit solver\n",
        "**autograd** is a Python package that calculates the gradients at the same time as the values by overloading the Numpy operations."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IkMiiSVY6VD9"
      },
      "outputs": [],
      "source": [
        "# this does not work as expected since least_squares is not \n",
        "# supported by autograd\n",
        "def spoly_interp_for_autograd(xs,ts,t):\n",
        "    return spoly_interp(xs,ts,t, pricing_helper)\n",
        "\n",
        "spi_grad = grad(spoly_interp_for_autograd)\n",
        "try:\n",
        "    print(spi_grad(ts,xs,t))\n",
        "except Exception as e:\n",
        "    print(f'Does not work, exception: {e}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "tDRlkplKBYFx"
      },
      "source": [
        "## Modify spoly_interp to calculate the gradients to the inputs xs using the naive Implicit Function Theorem"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "uuAWvJi8Bmrv"
      },
      "source": [
        "Extend PricingHelper to calculate the potential $\\frac{dw}{dx}$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-cBn4CEEHXzy"
      },
      "outputs": [],
      "source": [
        "class PricingHelperIft(PricingHelper):\n",
        "    '''\n",
        "    We simulate a typical programming pattern where auxilliary variables such \n",
        "    as w come wrapped in various helpers, etc. This is not strictly necessary\n",
        "    in this code but will be used later to illustrate some points.\n",
        "    '''\n",
        "    def __init__(self, w):\n",
        "        super().__init__(w)\n",
        "    \n",
        "    def update(self, xs, ts):\n",
        "        super().update(xs,ts)\n",
        "\n",
        "        # Capture the gradients if w is in fact a function of x. We could call\n",
        "        # autograd here but choose to code this by hand for brevity.\n",
        "        if self.updatable:\n",
        "            self.dw_dx_ = np.vstack((np.zeros_like(xs), 2*xs))\n",
        "        else:\n",
        "            self.dw_dx_ = np.zeros((2,len(xs)))\n",
        "     "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "j0Z-Ze2kGVyP"
      },
      "source": [
        "Modify **spoly_inter(...)** calling autograd when needed and implementing the IFT logic manually.\n",
        "\n",
        "The variable **c_fit** corresponds to $C(x,W(x))$ in the paper. Note that this driver should be aware of the variable $w$ to calculate **dobj_dw** i.e. $\\frac{\\partial \\Omega}{\\partial w}$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "fLmgOEXE6VD-"
      },
      "outputs": [],
      "source": [
        "def spoly_interp_ift(xs, ts, t, pricing_helper):\n",
        "    '''\n",
        "    This is a modification of spoly_interp() that supports gradients via \n",
        "    Naive IFT. We use autograd and need to use manual gradient manipulations\n",
        "    to collect them all.\n",
        "    \n",
        "    The original function spoly_interp(...) best-fits a stretched polynomial to \n",
        "    (ts,xs) and evaluates it at t. Here pricing_helper (via pricing_helper.w_)\n",
        "    is defining the interpolation between knots.\n",
        "    '''    \n",
        "\n",
        "    # Update the weights w and extract the relevant gradients\n",
        "    pricing_helper.update(xs,ts)\n",
        "    dw_dx = pricing_helper.dw_dx_\n",
        "\n",
        "    # The original objective function\n",
        "    def obj_f(c, x, pricing_helper = pricing_helper):\n",
        "        return np.concatenate((pricing_helper.spolyval(c, ts) - x, pricing_helper.spolyval_derivative(c,ts)))\n",
        "\n",
        "    # We need an unwrapped version of the objective function for autograd \n",
        "    # to be able to calculate dobj_dw below.\n",
        "    def obj_f_wrapper(c, x, w):\n",
        "        helper_ = PricingHelper(w)\n",
        "        return np.concatenate((helper_.spolyval(c, ts) - x, helper_.spolyval_derivative(c,ts)))\n",
        "\n",
        "    x0 = np.zeros_like(ts)\n",
        "    res = least_squares(lambda c: obj_f(c,xs), x0)\n",
        "    c_fit = res.x\n",
        "    v = pricing_helper.spolyval(c_fit, t)\n",
        "    # calc the gradients using IFT\n",
        "    dobj_dc = jacobian(obj_f, argnum = 0)(c_fit,xs)\n",
        "    dobj_dx = jacobian(obj_f, argnum = 1)(c_fit,xs)\n",
        "    dc_dx = -np.linalg.lstsq(dobj_dc,dobj_dx, rcond = None)[0]\n",
        "\n",
        "    # Calculate the gradient with respect to w. We need to keep adding\n",
        "    # these for all \"hidden\" variables that are used in obj_f\n",
        "    w = np.array(pricing_helper.w_.copy()) # a bit of a hoop here for autograd\n",
        "    dobj_dw = jacobian(obj_f_wrapper, argnum = 2)(c_fit,xs,w)   \n",
        "    \n",
        "    dc_dw = -np.linalg.lstsq(dobj_dc,dobj_dw, rcond = None)[0]\n",
        "    dc_dx += (dc_dw @ dw_dx)\n",
        "\n",
        "    dv_dc = grad(spolyval, argnum = 0)(c_fit, t, w)\n",
        "    dv_dx = dv_dc @ dc_dx\n",
        "\n",
        "    # need to add the dw_dx contribution to the final valuation as well\n",
        "    dv_dw = grad(spolyval, argnum = 2)(c_fit, t, w)\n",
        "    dv_dx += dv_dw @ dw_dx\n",
        "\n",
        "    return v, dv_dx\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HXZBP_sJ6VEB"
      },
      "source": [
        "Calculate the gradients using naive IFT and compare to gradients by bumping calculated previously.\n",
        "\n",
        "The two are now expected to be somewhat different because the IFT application is an approximation per Section 5.2\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "AMVV_9MN6VEC"
      },
      "outputs": [],
      "source": [
        "pricing_helper = PricingHelperIft(w_to_use)\n",
        "v_ift, grad_ift = spoly_interp_ift(xs,ts,t,pricing_helper)\n",
        "print(f'value = {v_ift}')\n",
        "print(f'gradients by ift = {grad_ift}')\n",
        "print(f'gradients by bmp = {grad_bump}')\n",
        "print(f'difference in gradients = {grad_ift - grad_bump}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "JeQlX1z3CTXL"
      },
      "source": [
        "## Calculate the gradients using AAD + Automatic IFT"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cN2eMQ8gCV-0"
      },
      "source": [
        "We implement the adjoints in **PricingHelper**. In a true AAD library these are generated automatically."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zxr0u8m3IREz"
      },
      "outputs": [],
      "source": [
        "class PricingHelperAdj(PricingHelperIft):\n",
        "      def __init__(self, w):\n",
        "        super().__init__(w)\n",
        "\n",
        "      def spolyval_adj(self, c, ts, state_adj):\n",
        "        '''\n",
        "        Propagate the adjoints through spolyval. Normally generated \n",
        "        automatically by the AAD library.\n",
        "        '''\n",
        "        # make sure we accept a single float not just arrays\n",
        "        ts = np.atleast_1d(ts)\n",
        "        w=self.w_\n",
        "        nc = len(c)\n",
        "        nw = len(w)\n",
        "        nt = len(ts)\n",
        "\n",
        "        # Just like in spolyval\n",
        "        tsw = w[0]*ts + w[1]*ts**2\n",
        "\n",
        "        sp_bar = state_adj['sp_bar']\n",
        "\n",
        "        # the length of sp_bar changes depending on the number of outputs \n",
        "        # of spolyval which is given by nt, make sure we line up with the \n",
        "        # state_adj here\n",
        "        if len(sp_bar) != nt:\n",
        "          raise(f'sp_bar length {len(sp_bar)} is not equal to THE expected {nt}')\n",
        "\n",
        "        # Start the adjoints with whatever is on the state_adj already -- \n",
        "        # this is important\n",
        "        c_bar = state_adj['c_bar']\n",
        "        w_bar = state_adj['w_bar']\n",
        "\n",
        "        # Loop over the length of the output of spolyval\n",
        "        for i in range(nt):\n",
        "\n",
        "          for n in range(nc):\n",
        "            # accumulate adjoints to coefs\n",
        "            c_bar[n] += tsw[i]**n * sp_bar[i]\n",
        "            \n",
        "            # Zero-order term has no sensitivity to w's\n",
        "            if n==0: \n",
        "              continue\n",
        "\n",
        "            # accumulate adjoints for w's\n",
        "            w_bar[0] += c[n] * n * tsw[i]**(n-1) * ts[i] * sp_bar[i]\n",
        "            w_bar[1] += c[n] * n * tsw[i]**(n-1) * ts[i]**2 * sp_bar[i]\n",
        "\n",
        "\n",
        "        # put adjoints back in the state_adj\n",
        "        state_adj['c_bar'] = c_bar\n",
        "        state_adj['w_bar'] = w_bar\n",
        "\n",
        "      def spolyval_der_adj(self, c, ts, state_adj):\n",
        "        '''\n",
        "        Propagate the adjoints through spolyval_der. Normally generated \n",
        "        automatically by the AAD library.\n",
        "        '''\n",
        "        # make sure we accept a single float not just arrays\n",
        "        ts = np.atleast_1d(ts)\n",
        "        w=self.w_\n",
        "        nc = len(c)\n",
        "        nw = len(w)\n",
        "        nt = len(ts)\n",
        "\n",
        "        # Just like in spolyval_der\n",
        "        tsw = w[0]*ts + w[1]*ts**2\n",
        "        tsw_der = w[0] + w[1]*ts*2\n",
        "\n",
        "        sp_der_bar = state_adj['sp_der_bar'] * regularization_weight\n",
        "\n",
        "        # the length of sp_bar changes depending on the number of outputs \n",
        "        # of spolyval which is given by nt, make sure we line up with the\n",
        "        # state_adj here\n",
        "        if len(sp_der_bar) != nt:\n",
        "          raise(f'sp_bar length {len(sp_bar)} is not equal to THE expected {nt}')\n",
        "\n",
        "        # Start the adjoints with whatever is on the state_adj already -- \n",
        "        # this is important\n",
        "        c_bar = state_adj['c_bar']\n",
        "        w_bar = state_adj['w_bar']\n",
        "\n",
        "        # Loop over the length of the output of spolyval\n",
        "        for i in range(nt):\n",
        "\n",
        "          for n in range(nc):\n",
        "            # accumulate adjoints to coefs\n",
        "            if (n>0):\n",
        "               c_bar[n] += tsw[i]**(n-1) * n * tsw_der[i] * sp_der_bar[i]\n",
        "            \n",
        "            # Zero-order term has no sensitivity to w's\n",
        "            if n==0: \n",
        "              continue\n",
        "\n",
        "            # accumulate adjoints for w's\n",
        "            if (n>0):\n",
        "              w_bar[0] += c[n] * n * (  tsw[i]**(n-1) ) * sp_der_bar[i]\n",
        "              w_bar[1] += c[n] * n * (  tsw[i]**(n-1) * ts[i]*2) * sp_der_bar[i]\n",
        "            if (n>1):\n",
        "              w_bar[0] += c[n] * n * ( (n-1) * tsw[i]**(n-2) * ts[i] *tsw_der[i] ) * sp_der_bar[i]\n",
        "              w_bar[1] += c[n] * n * ( (n-1) * tsw[i]**(n-2) * ts[i]**2 * tsw_der[i] ) * sp_der_bar[i]\n",
        "\n",
        "\n",
        "        # put adjoints back in the state_adj\n",
        "        state_adj['c_bar'] = c_bar\n",
        "        state_adj['w_bar'] = w_bar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sbsfIWT8ETE9"
      },
      "source": [
        "Initialize the state for the adjoints."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pIAvFwSFznDs"
      },
      "outputs": [],
      "source": [
        "def init_state_adj(ncoefs):\n",
        "  '''\n",
        "  Initialize state_adj. This will be done by the AAD library.\n",
        "  '''\n",
        "  state_adj = {\n",
        "      'sp_bar': np.array([1]),\n",
        "      'sp_der_bar': np.zeros(ncoefs),    \n",
        "      'c_bar' : np.zeros(ncoefs),\n",
        "      'x_bar': np.zeros(ncoefs),\n",
        "      'w_bar' : np.zeros(2),\n",
        "      'f_bar': np.zeros(ncoefs),\n",
        "      }\n",
        "\n",
        "  return state_adj \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gNSoqSAeEeXZ"
      },
      "source": [
        "Adjoints for the objective function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hKkOTqHKDtei"
      },
      "outputs": [],
      "source": [
        "def obj_f_adj(c, ts, x, helper, state_adj):\n",
        "  '''\n",
        "  Propagate adjoints through obj_f -- done by the AAD library\n",
        "  '''\n",
        "  x_bar = state_adj['x_bar']\n",
        "  f_bar = state_adj['f_bar']\n",
        "  state_adj['sp_bar'], state_adj['sp_der_bar'] = np.split(f_bar, 2)\n",
        "  f_bar = state_adj['sp_bar']\n",
        "  x_bar -= f_bar\n",
        "  state_adj['x_bar'] = x_bar\n",
        "  helper.spolyval_adj(c, ts, state_adj)\n",
        "  helper.spolyval_der_adj(c, ts, state_adj)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z_1lFEtYmt9a"
      },
      "source": [
        "The main part, run **spoly_interp(...)** with AAD + AIFT.\n",
        "\n",
        "  This is a modification of **spoly_interp()** that supports gradients via AAD + AIFT.\n",
        "    Note that all the adjoint steps can be automatically derived from the valuation steps by the AAD library and there are no explicit gradient manipulations.    \n",
        "    The original function **spoly_interp(...)** best-fits the stretched polynomial to **(ts,xs)** and evaluates it at **t**.\n",
        "    Here **pricing_helper** (via **pricing_helper.w_**) is defining the interpolation between knots. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "11CDo3fnuiw1"
      },
      "outputs": [],
      "source": [
        "def spoly_interp_aad(xs, ts, t, pricing_helper):\n",
        "    # Step 0. Initialize the state_adj\n",
        "    ncoefs = len(ts)\n",
        "    state_adj = init_state_adj(ncoefs)\n",
        "\n",
        "    # Step 1. Update the weights w and extract the relevant gradients\n",
        "    pricing_helper.update(xs,ts)\n",
        "    dw_dx = pricing_helper.dw_dx_\n",
        "\n",
        "\n",
        "    # Need to (re)define obj_f to take both coefs and xs as arguments so we can differentiate\n",
        "    def obj_f(c, x, pricing_helper=pricing_helper):\n",
        "        #return (helper_.spolyval(c, ts) - x)\n",
        "        return np.concatenate( (pricing_helper.spolyval(c, ts) - x, pricing_helper.spolyval_derivative(c,ts) ) )\n",
        "    \n",
        "    \n",
        "    # Step 2. Fit the objective function and extract the coefs c we fit\n",
        "    x0 = np.zeros_like(ts)\n",
        "    res = least_squares(lambda c: obj_f(c,xs), x0)\n",
        "    c_fit = res.x\n",
        "    \n",
        "    # Step 3. Calculate the value of spolyfit using the fitted coefs c_fit\n",
        "    v = pricing_helper.spolyval(c_fit, t)\n",
        "    print('val', v)\n",
        "\n",
        "    # Gradient to coefs, known in the Newton method so no extra calcs here\n",
        "    dobj_dc = jacobian(obj_f, argnum = 0)(c_fit, xs, pricing_helper)\n",
        "\n",
        "    # Adjoint for Step 3. I.e. propagate backwards until the call to the solver\n",
        "    pricing_helper.spolyval_adj(c_fit, t, state_adj)\n",
        "    c_bar = state_adj['c_bar']\n",
        "    \n",
        "    # Compute the correct adjoints of the objective function:\n",
        "    obj_f_bar = -np.linalg.lstsq(dobj_dc.T, c_bar, rcond = None)[0]\n",
        "    state_adj['f_bar'] = obj_f_bar\n",
        "\n",
        "    # print('obj_f_bar',obj_f_bar)\n",
        "    \n",
        "    # Adjoint for Step 2. Propagate through the objective function. Note that\n",
        "    # we do not have to compute dobj_dw unlike the Naive IFT approach\n",
        "    obj_f_adj(c_fit, ts, xs, pricing_helper, state_adj)\n",
        "    x_bar = state_adj['x_bar']\n",
        "    w_bar = state_adj['w_bar']\n",
        "    \n",
        "    # Adjoint for Step 1. Propagate through w=w(x)\n",
        "    x_bar += w_bar @ dw_dx\n",
        "\n",
        "    return v, x_bar"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "_bCRB8mfHcPZ"
      },
      "source": [
        "Calculate the gradients using AAD +AIFT and compare to the already computed gradients by bumping."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ZcDVrh-eHdDM"
      },
      "outputs": [],
      "source": [
        "pricing_helper = PricingHelperAdj(w_to_use)\n",
        "v_aad, grad_aad = spoly_interp_aad(xs,ts,t,pricing_helper)\n",
        "print(f'value = {v_aad}')\n",
        "print(f'gradients by aad = {grad_aad}')\n",
        "print(f'gradients by bmp = {grad_bump}')\n",
        "print(f'difference in gradients = {grad_aad - grad_bump}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## The end"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "LM_AIFT.ipynb",
      "provenance": []
    },
    "interpreter": {
      "hash": "c75fd95ee6a8673db96eeaeea4fa8e27c3cb071aef74affedb213b23408cb297"
    },
    "kernelspec": {
      "display_name": "Python 3.8.8 64-bit ('base': venv)",
      "name": "python3"
    },
    "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.8.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}