{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "committed-cincinnati",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "accepting-valentine",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1.0000, 1.0000, 1.0000, 1.0000, 0.5000], requires_grad=True)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#set states\n",
    "stocks = torch.ones(5)\n",
    "#Last one is different\n",
    "stocks[-1] = 0.5\n",
    "#now add the tracking requirement in place\n",
    "stocks.requires_grad_()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "reserved-parallel",
   "metadata": {},
   "outputs": [],
   "source": [
    "launch = torch.ones(5, requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "expected-consensus",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "debris = torch.tensor([2.2],requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "cross-drain",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Parameters\n",
    "SCALING = torch.ones(5)\n",
    "DELTA = 0.9 \n",
    "launch_debris_rate = 0.005\n",
    "collision_debris_rate = 0.0007"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "charitable-frost",
   "metadata": {},
   "outputs": [],
   "source": [
    "def survival(stock, debris):\n",
    "    #Gompertz distribution for simplicity\n",
    "    #commonly used with saturation\n",
    "    #TODO: ACTUALLY DERIVE A SURVIVAL FUNCTION. THIS IS JUST A PLACEHOLDER. PROBABLY SHOULD BE AN EXPONENTIAL DISTRIBUTION\n",
    "    eta = 1.0/(SCALING@stock)\n",
    "    b = 1/debris\n",
    "    \n",
    "    return 1 - ( b*eta*torch.exp(eta+b*stock-eta*torch.exp(b*stock)))\n",
    "\n",
    "\n",
    "def g(stock, debris, launches):\n",
    "    \n",
    "    new_stock = stock*survival(stock,debris) + launches\n",
    "    \n",
    "    #TODO: Currently Ignoring autocatalysis\n",
    "    new_debris = (1-DELTA)*debris + launch_debris_rate * launches.sum() + collision_debris_rate*(1-survival(stock,debris)) @ stock\n",
    "    \n",
    "    return (new_stock, new_debris)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "nonprofit-vintage",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "((tensor([[0.8457, 0.0271, 0.0271, 0.0271, 0.0271],\n",
       "          [0.0271, 0.8457, 0.0271, 0.0271, 0.0271],\n",
       "          [0.0271, 0.0271, 0.8457, 0.0271, 0.0271],\n",
       "          [0.0271, 0.0271, 0.0271, 0.8457, 0.0271],\n",
       "          [0.0126, 0.0126, 0.0126, 0.0126, 0.8731]]),\n",
       "  tensor([[0.0825],\n",
       "          [0.0825],\n",
       "          [0.0825],\n",
       "          [0.0825],\n",
       "          [0.0317]]),\n",
       "  tensor([[1., 0., 0., 0., 0.],\n",
       "          [0., 1., 0., 0., 0.],\n",
       "          [0., 0., 1., 0., 0.],\n",
       "          [0., 0., 0., 1., 0.],\n",
       "          [0., 0., 0., 0., 1.]])),\n",
       " (tensor([[4.2202e-05, 4.2202e-05, 4.2202e-05, 4.2202e-05, 1.2820e-05]]),\n",
       "  tensor([[0.0997]]),\n",
       "  tensor([[0.0050, 0.0050, 0.0050, 0.0050, 0.0050]])))"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Examples\n",
    "survival(stocks,debris)\n",
    "\n",
    "#Get the derivatives seperately\n",
    "jacobian(survival, (stocks,debris))\n",
    "\n",
    "#Get the derivatives as a single result\n",
    "torch.cat(jacobian(survival, (stocks,debris)), axis=1)\n",
    "\n",
    "#Testing state updates\n",
    "g(stocks, debris, launch)\n",
    "\n",
    "#Note the two tuples of jacobians: the first is for stock evolution, the second is for debris evolution\n",
    "jacobian(g, (stocks,debris,launch))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "future-greenhouse",
   "metadata": {},
   "source": [
    "## Next step: Construct the intertemporal-transition function(s)\n",
    " - Note: There are a couple of different ways to do this\n",
    "     - Just a single period transition function, manually iterated\n",
    "     - A recursive function that creates a $p$ period iterated function\n",
    "     - A recursive function that returns a list of functions iterated from 1 to $p$ periods\n",
    "\n",
    "I am planning on doing the latter, as each version is needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "higher-brook",
   "metadata": {},
   "outputs": [],
   "source": [
    "#setup\n",
    "beta = torch.tensor([0.95])\n",
    "util_weights = torch.tensor([1.0,1.0,0,0,0])\n",
    "sigma = 0.5\n",
    "\n",
    "w = torch.zeros(5)\n",
    "w[0]=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "approximate-motivation",
   "metadata": {},
   "outputs": [],
   "source": [
    "#This is not a good specification of the profit function, but it will work for now.\n",
    "def util(x):\n",
    "    return (util_weights@x)\n",
    "\n",
    "def single_transition(w, g, util, stocks, debris, launch, beta):\n",
    "    #TODO: change launch from a direct tensor, to a function.\n",
    "    bA = beta * jacobian(g, (stocks,debris,launch))[0][0]\n",
    "    return  bA.inverse() @ (w - jacobian(util,stocks))\n",
    "\n",
    "#need to create recursive transitions\n",
    "def recurse_transitions(w, g, util, stocks, debris, launch, beta, iterations):\n",
    "    #This is of type two from the discussion above\n",
    "    if iterations <= 1:\n",
    "        return single_transition(w, g, util, stocks, debris, launch, beta)\n",
    "    else: \n",
    "        #Get recursive results\n",
    "        curse = recurse_transitions(w,g,util,stocks,debris,launch,beta,iterations-1)\n",
    "        \n",
    "        #Get updated stocks and debris \n",
    "        stocks_iterated,debris_iterated = g(stocks, debris, launch)        \n",
    "        \n",
    "        #Return the updated values\n",
    "        return single_transition(curse, g, util, stocks_iterated, debris_iterated, launch, beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "prescribed-unemployment",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0372, -1.2487,  0.0372,  0.0372,  0.0164])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "single_transition(w,g,util,stocks,debris,launch,beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ordinary-admission",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.0372, -1.2487,  0.0372,  0.0372,  0.0164])\n",
      "tensor([-1.0135, -2.3671,  0.0391,  0.0391,  0.0172])\n",
      "tensor([-2.1195, -3.5443,  0.0412,  0.0412,  0.0181])\n"
     ]
    }
   ],
   "source": [
    "r1 = recurse_transitions(w, g, util, stocks, debris, launch, beta, 1)\n",
    "r2 = recurse_transitions(w, g, util, stocks, debris, launch, beta, 2)\n",
    "r3 = recurse_transitions(w, g, util, stocks, debris, launch, beta, 3)\n",
    "\n",
    "print(r1)\n",
    "print(r2)\n",
    "print(r3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "committed-porcelain",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([1.8600, 1.8600, 1.8600, 1.8600, 1.4401], grad_fn=<AddBackward0>),\n",
       " tensor([0.2454], grad_fn=<AddBackward0>),\n",
       " tensor([2.8600, 2.8600, 2.8600, 2.8600, 2.4401], grad_fn=<AddBackward0>),\n",
       " tensor([0.0495], grad_fn=<AddBackward0>))"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#TODO: manually check single_transition\n",
    "stocks_1,debris_1 = g(stocks, debris, launch) \n",
    "stocks_2,debris_2 = g(stocks_1, debris_1, launch)\n",
    "stocks_1,debris_1,stocks_2,debris_2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "touched-judge",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([ 0.0372, -1.2487,  0.0372,  0.0372,  0.0164])\n",
      "tensor([-1.0135, -2.3671,  0.0391,  0.0391,  0.0172])\n",
      "tensor([-2.1195, -3.5443,  0.0412,  0.0412,  0.0181])\n"
     ]
    }
   ],
   "source": [
    "#Recurse 1\n",
    "sin1 = single_transition(w,g,util,stocks,debris,launch,beta)\n",
    "#Recurse 2\n",
    "sin2 = single_transition(sin1 ,g,util,stocks_1,debris_1,launch,beta)\n",
    "#recurse 3\n",
    "sin3 = single_transition(sin2,g,util,stocks_2,debris_2,launch,beta)\n",
    "\n",
    "print(sin1)\n",
    "print(sin2)\n",
    "print(sin3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "cooked-absence",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-2.1195, -3.5443,  0.0412,  0.0412,  0.0181])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recurse_transitions(sin2, g, util, stocks_2, debris_2, launch, beta, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "accepting-grade",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-2.1195, -3.5443,  0.0412,  0.0412,  0.0181])"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#I think this highlights the error.\n",
    "recurse_transitions(recurse_transitions(sin1, g, util, stocks_1, debris_1, launch, beta, 1), g, util, stocks_1, debris_1, launch, beta, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "indoor-papua",
   "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}