{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Differentiation"
   ],
   "id": "site-title"
  },
  {
   "cell_type": "markdown",
   "id": "096c3f75",
   "metadata": {},
   "source": [
    "## Before you begin\n",
    "\n",
    "If you're new to calculus or need a refresher of the concepts and formulas, check out the Khan Academy course on YouTube. \n",
    "\n",
    "LINK : https://youtube.com/playlist?list=PL19E79A0638C8D449\n",
    "\n",
    "Specially for Derivatives: https://youtu.be/rAof9Ld5sOg "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72075517",
   "metadata": {},
   "source": [
    "## Simple differentiation\n",
    "\n",
    "Say we have this function: $f(x) = x^{4} - 5x^{3} + 2x^{2} +4x - 1$\n",
    "\n",
    "And we want to differentiate it in terms of $x$ or find its derivative which is $\\frac{df}{dx}$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "dc35e499",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.007334Z",
     "start_time": "2023-02-08T00:59:58.780595Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax as J"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "d72d8064",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.010805Z",
     "start_time": "2023-02-08T00:59:59.008903Z"
    }
   },
   "outputs": [],
   "source": [
    "def f(x):\n",
    "    return x**4 - 5 * x**3 + 2 * x**2 + 4 * x - 1"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98bd3d24",
   "metadata": {},
   "source": [
    "Why function? Good question. Jax was designed with Functional Programming in mind hence this function stuff. You can also use lambdas. For example this can also be written as \n",
    "\n",
    "```python\n",
    "f = lambda x: x**4 - 5 * x**3 + 2 * x**2 + 4 * x - 1\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "506fd90b",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.027126Z",
     "start_time": "2023-02-08T00:59:59.011625Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<function __main__.f(x)>"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_dx = J.grad(f)\n",
    "df_dx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb62303f",
   "metadata": {},
   "source": [
    "## Derivative Where?"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc25777b",
   "metadata": {},
   "source": [
    "Where's the derivative? Huh Charlston! Where's the derivative? \n",
    "\n",
    "Well you didn't tell me for which value of x you need a derivative of y here!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "85dc4cac",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.418782Z",
     "start_time": "2023-02-08T00:59:59.027960Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-16., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_dx(2.) # make sure to pass in floats, that's how JAX likes it!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a8f87d6",
   "metadata": {},
   "source": [
    "Let's verify that, shall we? \n",
    "\n",
    "$\\frac{df}{dx} = 4x^{3} - 15x^{2} + 4x + 4$\n",
    "\n",
    "So for $x = 2.0$, or say $x = 2$\n",
    "\n",
    "$\\frac{df}{dx} = 4x^{3} - 15x^{2} + 4x + 4 = 4 \\times 2^{3} - 15 \\times 2^{2} + 4 \\times 2 + 4 = -16$\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b56a3977",
   "metadata": {},
   "source": [
    "## Higher order differentiation\n",
    "\n",
    "What we saw above was the first order derivative of  y in terms of x. So a very simple definition would be that when you differentiate once, it's a first order derivative, twice - second order derivative and so on. Anything above first order derivative is called a higher order derivative. \n",
    "\n",
    "Let's try the second and third order derivative:\n",
    "\n",
    "$\\frac{d}{dx}(\\frac{df}{dx}) = \\frac{d}{dx}(4x^{3} - 15x^{2} + 4x + 4) = 12x^{2} - 30x + 4$\n",
    "\n",
    "This second order derivative or, $\\frac{d}{dx}(\\frac{df}{dx})$ can also be written as $\\frac{d^{2}f}{dx^{2}} $ or simply ${f}''$. Guess what the first order derivative is then? ${f}'$. \n",
    "\n",
    "So, \n",
    "\n",
    "${f}''' = 24x - 30$\n",
    "\n",
    "Fun fact, for a fourth order derivative you don't have to add like 4 ' signs. You can replace ' with numbers as well, and just write something like $f_{4}$ or $f^{4}$ or $f^{iv}$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "425efc42",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.540717Z",
     "start_time": "2023-02-08T00:59:59.419709Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-8., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f2 = J.grad(df_dx)\n",
    "f2(2.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "fdc68abd",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.644427Z",
     "start_time": "2023-02-08T00:59:59.541576Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-6., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f3 = J.grad(f2)\n",
    "f3(1.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "42711e35",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.763268Z",
     "start_time": "2023-02-08T00:59:59.645349Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(24., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f4 = J.grad(f3)\n",
    "f4(0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "6b3fff9d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.904031Z",
     "start_time": "2023-02-08T00:59:59.764188Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f5 = J.grad(f4)\n",
    "f5(1.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80247abe",
   "metadata": {},
   "source": [
    "Why's this 0? Like all those books out there, this has been left as an exercise for the reader. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6fc4bee",
   "metadata": {},
   "source": [
    "## Partial Derivatives And Multvariate Differentiation\n",
    "\n",
    "So far the derivatives have been in terms of only one variable, $x$. What if you have multiple variables? For example:\n",
    "$f(x, y, z) = 3x + 2z + 3xy + 5zx$\n",
    "\n",
    "The partial derivatives of $f$ here can be:\n",
    "\n",
    "$\\frac{\\partial f}{\\partial x}$, $\\frac{\\partial f}{\\partial y}$, $\\frac{\\partial f}{\\partial z}$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c316d672-3748-4b72-8805-931b166007ba",
   "metadata": {},
   "source": [
    "--------\n",
    "**Note:**  \n",
    "The function defined above in the code, f(x), had only one parameter x. By default JAX takes derivatives in term of the first such parameter. But you can get around that by using argnums. \n",
    "\n",
    "------------------------------"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a1890515",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.906931Z",
     "start_time": "2023-02-08T00:59:59.905096Z"
    }
   },
   "outputs": [],
   "source": [
    "# define the function f\n",
    "def f(x, y, z):\n",
    "    return 3*x + 2*z + 3*x* y + 5*z*x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "987b52da",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T00:59:59.955339Z",
     "start_time": "2023-02-08T00:59:59.907820Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(3., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_y = J.grad(f, argnums=1) # in terms of the second param, y\n",
    "f1_y(1., 2., 3.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "8209bd21",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:00:00.003619Z",
     "start_time": "2023-02-08T00:59:59.957430Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(7., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "f1_z = J.grad(f, argnums=2) # in terms of z\n",
    "f1_z(1., 2., 3.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1343026d",
   "metadata": {},
   "source": [
    "What about $x$ ? Do we leave it hanging there? Okay....."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "09efec56",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:00:00.011437Z",
     "start_time": "2023-02-08T01:00:00.004678Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(24., dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# by default argnums is 0 so even if you don't \n",
    "# mention anything it will use x\n",
    "\n",
    "f1_x = J.grad(f, argnums=0)\n",
    "f1_x(1., 2., 3.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d896478",
   "metadata": {},
   "source": [
    "## Chain Rule\n",
    "\n",
    "So there can be multiple variables. What if they're linked? Like one depends on the other? Like a chain? \n",
    "\n",
    "![godzilla.jpg](../images/godzilla.jpg)\n",
    "\n",
    "\n",
    "Actually it's not that difficult. You just have to understand the how some variables can be dependent on each other. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11532225",
   "metadata": {},
   "source": [
    "Let's think in terms of functions, since JAX looooves functions. \n",
    "\n",
    "Let's say I want to cook some noodles. Nothing fancy, just some good noodles boiled up, tossed into a pan with the ingredients of your choice.\n",
    "\n",
    "\n",
    "So first you have to $boil()$ the noodle. Then you need to prepare the pan with oil and $fry()$ the other ingredients, then you need to pray and throw in all the noodle on the pan. \n",
    "\n",
    "So in a functional way, you can say it like this: `fry(boil(noodles), some stuff, pray)`\n",
    "\n",
    "But wait, that's a function inside a function! Oh no! But anyway. These kind of functions which can hold other functions are called composite functions ([more on this here](https://www.khanacademy.org/math/ap-calculus-ab/ab-differentiation-2-new/ab-3-1a/a/chain-rule-review)). And when you try to differentiate them, you need the chain rule. "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "891d93c3",
   "metadata": {},
   "source": [
    "## Why though?\n",
    "\n",
    "![why](../images/giphy.gif)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e6d591",
   "metadata": {},
   "source": [
    "Differentiation has a simple goal : to find out how much effect can the tiniest of the changes can make to something. Making cookies? How much change in sugar can change the taste? Baking potatoes? How slight change of temperature can ruin them for you? Actually real life problems regarding this are far more intricate but let's not get into those. \n",
    "\n",
    "So when you get composite functions which depend on each other, like you can't fry without boiling the noodles first, you need to use the chain rule to differentiate and observe the effects. Enough chit chat and bad examples, we'll dive into code now."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "86b9a9e7",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:00:00.014130Z",
     "start_time": "2023-02-08T01:00:00.012452Z"
    }
   },
   "outputs": [],
   "source": [
    "import jax.numpy as jnp\n",
    "\n",
    "def some_composite_function(x):\n",
    "    return jnp.log(\n",
    "        jnp.sin(\n",
    "            jnp.cos(x)\n",
    "        )\n",
    "    )\n",
    "\n",
    "# from a naive point of view, if you call composites nested functions, you may not be altogether wrong!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "5eeb422f",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:00:00.270534Z",
     "start_time": "2023-02-08T01:00:00.015046Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-1.4679605, dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d_composite = J.grad(some_composite_function)\n",
    "d_composite(45.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70f86541",
   "metadata": {},
   "source": [
    "Seems complicated? Let's break them up and then see how it goes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0864905d",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:00:00.274545Z",
     "start_time": "2023-02-08T01:00:00.271879Z"
    }
   },
   "outputs": [],
   "source": [
    "def a(x):\n",
    "    return jnp.log(x)\n",
    "\n",
    "def b(x):\n",
    "    return jnp.sin(x)\n",
    "\n",
    "def c(x):\n",
    "    return jnp.cos(x)\n",
    "\n",
    "def comp(x):\n",
    "    return a(b(c(x)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1edc0a84",
   "metadata": {},
   "source": [
    "So what's happening here is that a relies on b, b relies on output from c. And when you run grad,"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "127cf1c9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2023-02-08T01:00:00.291662Z",
     "start_time": "2023-02-08T01:00:00.275573Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-1.4679605, dtype=float32, weak_type=True)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "d_c = J.grad(comp)\n",
    "d_c(45.)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "22773a35",
   "metadata": {},
   "source": [
    "Where's the effect though? I understand that this isn't enough to give you feeling of where the effect lands actually, but if you keep your patience until we check backpropagation on a learning algorithm(stop calling these glorious learning functions AI, they're not!), you'll see it for yourself!\n",
    "\n",
    "## More on Autodiff\n",
    "\n",
    "I found [this video on YouTube to be a very friendly and useful introduction to Autodiff](https://youtu.be/wG_nF1awSSY)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "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.12.11"
  },
  "toc": {
   "base_numbering": 1,
   "nav_menu": {},
   "number_sections": true,
   "sideBar": true,
   "skip_h1_title": false,
   "title_cell": "Table of Contents",
   "title_sidebar": "Contents",
   "toc_cell": false,
   "toc_position": {},
   "toc_section_display": true,
   "toc_window_display": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}