{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "linear-harvey",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian\n",
    "import itertools"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "honey-excuse",
   "metadata": {},
   "source": [
    "# Setup Functions\n",
    "## General CompositionFunctions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "helpful-radical",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Set up functions to compose functions \n",
    "# These functions will \n",
    "#   - compose two functions together\n",
    "#   - compose a function to itself n times.\n",
    "\n",
    "def compose(f,g):\n",
    "    return lambda *args: f(g(*args))\n",
    "\n",
    "def compose_recursive_functions(fn,n):\n",
    "    #Set base conditions\n",
    "    out_func = None\n",
    "    out_func_list =[]\n",
    "\n",
    "    #build the compositions of functions\n",
    "    for f in itertools.repeat(fn, n):\n",
    "        if out_func == None:\n",
    "            out_func = f\n",
    "        else:\n",
    "            out_func = compose(f,out_func)\n",
    "\n",
    "        out_func_list.append(out_func)\n",
    "    \n",
    "    return out_func_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fifty-southwest",
   "metadata": {},
   "source": [
    "## Setup functions related to the problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "fancy-manual",
   "metadata": {},
   "outputs": [],
   "source": [
    "### Background functions\n",
    "\n",
    "\n",
    "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 laws_of_motion(stock, debris, launches):\n",
    "    \n",
    "    new_stock = stock*survival(stock,debris) + launches #TODO: Launches will become a function (neural network)\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",
    "\n",
    "#This is not a good specification of the profit function, but it will work for now.\n",
    "def profit(x):\n",
    "    return UTIL_WEIGHTS @ x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "dietary-vault",
   "metadata": {},
   "outputs": [],
   "source": [
    "def single_transition(item_to_iterate, laws_motion, profit,launch, stocks, debris ):\n",
    "    #TODO: change launch from a direct tensor, to a function.\n",
    "    \n",
    "    #Calculate the inverse\n",
    "    bA = BETA * jacobian(laws_motion, (stocks,debris,launch))[0][0]\n",
    "    #TODO: figure out some diagnostics for this section\n",
    "    \n",
    "    \n",
    "    return  bA.inverse() @ (item_to_iterate - jacobian(profit,stocks))\n",
    "\n",
    "# This function wraps the single transition and handles updating dates etc.\n",
    "def transition_wrapper(data_in):\n",
    "    #unpack states and functions\n",
    "    stocks, debris,profit, launch, laws_motion,item_to_transition = data_in\n",
    "    \n",
    "    #Calculate new states\n",
    "    new_stocks, new_debris = laws_motion(stocks,debris,launch)\n",
    "    \n",
    "    #WARNING: RECURSION: You may break your head...\n",
    "    #This gets the transition of the value function derivatives over time.\n",
    "    transitioned = single_transition(\n",
    "                    item_to_transition,  #item to iterate, i.e. the derivatives of the value function\n",
    "                    laws_motion, profit, launch, #functions #TODO: reimplement with launch as a function\n",
    "                    stocks, debris #states\n",
    "                    )\n",
    "    \n",
    "    #collects the data back together for return, including the updated state variables\n",
    "    data_out = new_stocks, new_debris, profit, launch, laws_motion, transitioned\n",
    "    \n",
    "    return data_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "illegal-thriller",
   "metadata": {},
   "source": [
    "# Actual calculations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "coated-dressing",
   "metadata": {},
   "outputs": [],
   "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_()\n",
    "\n",
    "#Setup Debris\n",
    "debris = torch.tensor([2.2],requires_grad=True)\n",
    "\n",
    "#CHANGE LATER: Launch is currently a value, should be a function (i.e. neural network)\n",
    "launch = torch.ones(5, requires_grad=True)\n",
    "\n",
    "#compose the functions together.\n",
    "base_data = (stocks,debris, profit, launch, laws_of_motion, torch.ones(5, requires_grad=True))\n",
    "\n",
    "#Parameters\n",
    "SCALING = torch.ones(5)\n",
    "DELTA = 0.9 \n",
    "LAUNCH_DEBRIS_RATE = 0.005\n",
    "COLLISION_DEBRIS_RATE = 0.0007\n",
    "UTIL_WEIGHTS = torch.tensor([1,-0.2,0,0,0])\n",
    "BETA = 0.95"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "analyzed-transfer",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor([-0.1543,  1.3888,  1.1316,  1.1316,  1.1553], grad_fn=<MvBackward>)\n",
      "tensor([-1.2150,  1.6724,  1.1912,  1.1912,  1.2161], grad_fn=<MvBackward>)\n",
      "tensor([-2.3316,  1.9710,  1.2539,  1.2539,  1.2801], grad_fn=<MvBackward>)\n",
      "tensor([nan, nan, nan, nan, nan], grad_fn=<MvBackward>)\n",
      "tensor([nan, nan, nan, nan, nan], grad_fn=<MvBackward>)\n"
     ]
    }
   ],
   "source": [
    "#calculate results for first 5 iterations\n",
    "for f in compose_recursive_functions(transition_wrapper,5):\n",
    "    result = f(base_data)\n",
    "    print(result[5])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "confidential-stadium",
   "metadata": {},
   "source": [
    "Note how this fails on the last few iterations.\n",
    "I need to get better model functions (profit, laws_of_motion, etc) together to test this out.\n",
    "Alternatively, I can check for areas where the determinant of $A$ is zero, possibly by doing some sort of grid search?\n",
    "\n",
    "Maybe with a standard RBC model?\n",
    "\n",
    "Also, maybe I can create a `Model` class that upon construction will capture the necesary constants, functions, etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbb2f9e4-2248-467e-b6f9-dd5cdd156ce6",
   "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
}