{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "11cc7082-6fb7-49fd-992e-6678f1570ed9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "9cff4845-9f88-4f49-8ad9-d94386ccc7dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "a = torch.tensor([2., 3.], requires_grad=True)\n",
    "b = torch.tensor([6., 4.], requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f49549e1-dbcb-4aa7-8036-8bef936b3b98",
   "metadata": {},
   "outputs": [],
   "source": [
    "def f(a,b):\n",
    "    return a**3+a**4 + a*b**2 -b**2 -b**5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7ee977aa-d71e-4cfd-8743-a06216fa2209",
   "metadata": {},
   "outputs": [],
   "source": [
    "f(a,b).backward(q)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "bccc63c6-ce7c-402e-99fc-d5b0a9b0deb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 92., 169.])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "a.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "fe0575a3-bdc9-4b2f-9564-b01763c09c2b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-6252., -1168.])"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "b.grad"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "e653ec77-4034-4797-b5db-bb06778090da",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Begin demo \n",
      "\n",
      "x = tensor([4.], requires_grad=True)\n",
      "y = tensor([29.], grad_fn=<AddBackward0>)\n",
      "\n",
      "df = <AddBackward0 object at 0x7f388c0cd640>\n",
      "\n",
      "gradient of func(x) = \n",
      "tensor([11.])\n"
     ]
    }
   ],
   "source": [
    "# gradient_demo.py\n",
    "\n",
    "import torch as T\n",
    "device = T.device(\"cpu\")\n",
    "\n",
    "def some_func(x):\n",
    "  result = (x * x) + (3 * x) + 1\n",
    "  return result\n",
    "\n",
    "def main():\n",
    "  print(\"\\nBegin demo \\n\")\n",
    "\n",
    "  x = T.tensor([4.0], dtype=T.float32,\n",
    "    requires_grad=True).to(device)\n",
    "  y = some_func(x)\n",
    "\n",
    "  print(\"x = \" + str(x))\n",
    "  print(\"y = \" + str(y))\n",
    "  print(\"\")\n",
    "\n",
    "  df = y.grad_fn\n",
    "  print(\"df = \" + str(df))\n",
    "  print(\"\")\n",
    "\n",
    "  y.backward()  # compute grad of some_func(4)\n",
    "  print(\"gradient of func(x) = \")\n",
    "  print(x.grad) # 2(4) + 3 = 11\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "  main()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3470450-834c-43db-8f89-d6f51e2754e0",
   "metadata": {},
   "source": [
    "# Try this\n",
    " - define a bellman like function \n",
    "     - use `def` \n",
    " - Find the derivatives with respect to X\n",
    " - Find the derivatives with respect to $\\theta$\n",
    "     - Invert this to find the time-transition function?\n",
    "     - Inversion is probably impossible. Instead, solve both the policy and value of derivatives together?\n",
    "     \n",
    "     \n",
    "## Math for conditions\n",
    "### Optimality\n",
    "Stationarity and Complementary Slackness conditions\n",
    "$$\n",
    "0 = \\frac{\\partial F}{\\partial x} + \\beta \\frac{\\partial G}{\\partial x} \\frac{\\partial V}{\\partial G} + \\lambda \\\\\n",
    "0 = \\lambda \\cdot g(x,\\theta) \\\n",
    "$$\n",
    "### Envelope\n",
    "$$\n",
    "0 = \\frac{\\partial F}{\\partial \\theta} + \\frac{\\partial G}{\\partial \\theta} \\frac{\\partial V}{\\partial G} - \\frac{\\partial V}{\\partial \\theta}\n",
    "$$"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91e8528c-a2c4-4dfd-87c6-3d367c9d0f2e",
   "metadata": {},
   "source": [
    "So, how do you incorporate the situation where you have to iterate multiple times?\n",
    " - Just add conditions as rows?\n",
    " - Solve and substitute using some theorem on the inverse of derivatives in multivariate systems?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9b9df983-d2b2-4515-a98b-2469d8d392f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def utility(c,a=2):\n",
    "    return (c**(1-a))/(1-a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "ea3a604f-48bb-424b-bb76-649d640a55db",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-1.0"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "8cc8b408-c0b9-4344-bed7-f6fbb6514a2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def bellman(theta, x, V, b=0.95):\n",
    "    pass\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7730732-de9e-4d8c-b85f-112bee09762a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.9.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}