{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "innovative-filename",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "filled-question",
   "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": "middle-lease",
   "metadata": {},
   "outputs": [],
   "source": [
    "launch = torch.ones(5, requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "regular-grounds",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "debris = torch.tensor([2.2],requires_grad=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "satisfied-hawaiian",
   "metadata": {},
   "outputs": [],
   "source": [
    "scaling = torch.ones(5)\n",
    "delta = 0.9 \n",
    "launch_debris_rate = 0.05\n",
    "collision_debris_rate = 0.07"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "saved-corporation",
   "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",
    "    new_stock = stock*survival(stock,debris) + launches\n",
    "    \n",
    "    #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": "nominated-visitor",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([0.8600, 0.8600, 0.8600, 0.8600, 0.8802], grad_fn=<RsubBackward1>)"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "survival(stocks,debris)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "clean-panama",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[-0.0142,  0.0271,  0.0271,  0.0271,  0.0271],\n",
       "         [ 0.0271, -0.0142,  0.0271,  0.0271,  0.0271],\n",
       "         [ 0.0271,  0.0271, -0.0142,  0.0271,  0.0271],\n",
       "         [ 0.0271,  0.0271,  0.0271, -0.0142,  0.0271],\n",
       "         [ 0.0251,  0.0251,  0.0251,  0.0251, -0.0142]]),\n",
       " tensor([[0.0825],\n",
       "         [0.0825],\n",
       "         [0.0825],\n",
       "         [0.0825],\n",
       "         [0.0634]]))"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Get the derivatives seperately\n",
    "jacobian(survival, (stocks,debris))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "cardiovascular-music",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.0142,  0.0271,  0.0271,  0.0271,  0.0271,  0.0825],\n",
       "        [ 0.0271, -0.0142,  0.0271,  0.0271,  0.0271,  0.0825],\n",
       "        [ 0.0271,  0.0271, -0.0142,  0.0271,  0.0271,  0.0825],\n",
       "        [ 0.0271,  0.0271,  0.0271, -0.0142,  0.0271,  0.0825],\n",
       "        [ 0.0251,  0.0251,  0.0251,  0.0251, -0.0142,  0.0634]])"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Get the derivatives as a single result\n",
    "torch.cat(jacobian(survival, (stocks,debris)), axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "funky-illness",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([1.8600, 1.8600, 1.8600, 1.8600, 1.4401], grad_fn=<AddBackward0>),\n",
       " tensor([0.5134], grad_fn=<AddBackward0>))"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Testing state updates\n",
    "g(stocks, debris, launch)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "charged-dairy",
   "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([[0.0042, 0.0042, 0.0042, 0.0042, 0.0013]]),\n",
       "  tensor([[0.0747]]),\n",
       "  tensor([[0.0500, 0.0500, 0.0500, 0.0500, 0.0500]])))"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#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": "marked-flower",
   "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": 12,
   "id": "outdoor-action",
   "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "unavailable-hawaii",
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "composite-cooperative",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2., grad_fn=<DotBackward>)"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "util(stocks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "activated-advancement",
   "metadata": {},
   "outputs": [],
   "source": [
    "w = torch.zeros(5)\n",
    "w[0]=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "wireless-concord",
   "metadata": {},
   "outputs": [],
   "source": [
    "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))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "detected-cooking",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0372, -1.2487,  0.0372,  0.0372,  0.0164])"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "single_transition(w,g,util,stocks,debris,launch,beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "sharp-pound",
   "metadata": {},
   "outputs": [],
   "source": [
    "#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": 19,
   "id": "radical-reason",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0372, -1.2487,  0.0372,  0.0372,  0.0164])"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recurse_transitions(w, g, util, stocks, debris, launch, beta, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "intelligent-angle",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.4740, -1.0278, -0.0434, -0.0434, -0.0575])"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recurse_transitions(w, g, util, stocks, debris, launch, beta, 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "jewish-lemon",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.7046, -0.9430, -0.0884, -0.0884, -0.1102])"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recurse_transitions(w, g, util, stocks, debris, launch, beta, 3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "scenic-literature",
   "metadata": {},
   "outputs": [],
   "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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "moving-slave",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([ 0.0372, -1.2487,  0.0372,  0.0372,  0.0164])"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Recurse 1\n",
    "sin1 = single_transition(w,g,util,stocks,debris,launch,beta)\n",
    "sin1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "corrected-radar",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.4740, -1.0278, -0.0434, -0.0434, -0.0575])"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#Recurse 2\n",
    "sin2 = single_transition(sin1 \n",
    "                         ,g\n",
    "                         ,util\n",
    "                         ,stocks_1\n",
    "                         ,debris_1\n",
    "                         ,launch,beta)\n",
    "sin2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "id": "handy-pencil",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1.5516, -2.1345, -0.0456, -0.0456, -0.1302])"
      ]
     },
     "execution_count": 36,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#recurse 3\n",
    "single_transition(sin2,g,util,stocks_2,debris_2,launch,beta)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "id": "attempted-affairs",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-1.5516, -2.1345, -0.0456, -0.0456, -0.1302])"
      ]
     },
     "execution_count": 37,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "recurse_transitions(sin2, g, util, stocks_2, debris_2, launch, beta, 1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "id": "honest-diana",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([-0.7046, -0.9430, -0.0884, -0.0884, -0.1102])"
      ]
     },
     "execution_count": 39,
     "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": "animated-nudist",
   "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
}