Orbits/Code/ImplementLoss.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "indie-evolution",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian\n",
    "import itertools\n",
    "import math\n",
    "\n",
    "class States():\n",
    "    \"\"\"\n",
    "    This is supposed to capture the state variables of the model\n",
    "    \"\"\"\n",
    "    def __init__(self, stocks,debris):\n",
    "        pass\n",
    "\n",
    "    def transition(self):\n",
    "        #tying the transition functions in here might be useful.\n",
    "        pass\n",
    "\n",
    "    \n",
    "class LDAPV():\n",
    "    \"\"\"\n",
    "    This class represents the outputs from the neural network.\n",
    "    They are of two general categories:\n",
    "        - the launch functions \n",
    "    \"\"\"\n",
    "    def __init__(self, launches, partials):\n",
    "        self.launches = launches\n",
    "        self.partials = partials\n",
    "    \n",
    "    def launch_single(constellation):\n",
    "        #returns the launch decision for the constellation of interest\n",
    "        \n",
    "        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.launches @ filter_tensor\n",
    "    \n",
    "    def launch_vector(constellation):\n",
    "        #returns the launch decision for the constellation of interest\n",
    "        \n",
    "        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.launches * filter_tensor\n",
    "    \n",
    "    def partial_vector(constellation):\n",
    "        #returns the partials of the value function corresponding to the constellation of interest\n",
    "        \n",
    "        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.partials @ filter_tensor\n",
    "    \n",
    "    def partial_matrix(constellation):\n",
    "        #returns the partials of the value function corresponding to the constellation of interest\n",
    "        \n",
    "        filter_tensor = torch.zeros(NUMBER_CONSTELLATIONS)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.partials * filter_tensor\n",
    "    \n",
    "    def __str__(self):\n",
    "        return \"Launch Decisions and Partial Derivativs of value function with\\n\\t states\\n\\t\\t {}\\n\\tPartials\\n\\t\\t{}\".format(self.states,self.partials)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "stuffed-firmware",
   "metadata": {},
   "source": [
    "# Setup Functions\n",
    "## General CompositionFunctions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "mexican-serial",
   "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",
    "    \"\"\"\n",
    "    this function composes two functions f and g in\n",
    "    the order f(g(*args))\n",
    "    \n",
    "    it returns a function\n",
    "    \"\"\"\n",
    "    return lambda *args: f(g(*args))\n",
    "\n",
    "def compose_recursive_functions(fn,n):\n",
    "    \"\"\"\n",
    "    This function takes a function fn and composes it with itself n times\n",
    "    \n",
    "    Returns an array/list that contains each of the n composition levels, ordered\n",
    "    from fn() to fn^n()\n",
    "    \"\"\"\n",
    "\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 first iteration\n",
    "        if out_func == None:\n",
    "            out_func = f\n",
    "        else:\n",
    "            out_func = compose(f,out_func)\n",
    "\n",
    "        #append the ou\n",
    "        out_func_list.append(out_func)\n",
    "    \n",
    "    return out_func_list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "public-alloy",
   "metadata": {},
   "source": [
    "# functions related to transitions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "advised-enemy",
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "invalid syntax (<ipython-input-3-2a8ca63b5912>, line 52)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  File \u001b[0;32m\"<ipython-input-3-2a8ca63b5912>\"\u001b[0;36m, line \u001b[0;32m52\u001b[0m\n\u001b[0;31m    launch = neural_network.forward().\u001b[0m\n\u001b[0m                                      ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m invalid syntax\n"
     ]
    }
   ],
   "source": [
    "def single_transition(laws_motion_fn, profit_fn, stocks, debris, neural_network):\n",
    "    \"\"\"\n",
    "    This function represents the inverted envelope conditions.\n",
    "    It allows us to describe the derivatives of the value function evaluated at time $t+1$ in terms based in time period $t$.\n",
    "    \n",
    "    It takes a few derivatives (jacobians), and due to the nature of the return as tuples\n",
    "    they must be concatenated.\n",
    "    \n",
    "    It returns the transitioned values\n",
    "    \"\"\"\n",
    "    \n",
    "    launch_and_partials = neural_network.forward(stocks,debris)\n",
    "    \n",
    "    #Get the jacobian\n",
    "    a = jacobian(laws_motion_fn, (stocks, debris, launch_and_partials.launches))\n",
    "    \n",
    "    #Reassemble the Jacobian nested tuples into the appropriate tensor\n",
    "    A = BETA * torch.cat((torch.cat((a[0][0],a[0][1]),dim=1),torch.cat((a[1][0],a[1][1]),dim=1)), dim=0)\n",
    "\n",
    "    #TODO: figure out some diagnostics for this section\n",
    "    #Possibilities include:\n",
    "    # - Det(A) ~= 0\n",
    "    # - EigVal(A) ~= 0\n",
    "    # - A.inverse() with a try catch system to record types of returns\n",
    "    #Alternatively, \n",
    "    #if abs(a.det())\n",
    "    \n",
    "    #Calculate the item to transition\n",
    "    f_jacobians = jacobian(profit_fn,(stocks, debris, launch_and_partials.launches))\n",
    "\n",
    "    #issue with shape here: my launch function is for all launches, not just a single launch.\n",
    "    f_theta = torch.cat([f_jacobians[0][0], f_jacobians[1][0]],axis=0) \n",
    "\n",
    "    #FIX: issue with scaling here, in that I need to get the constellation level partials, not the full system.\n",
    "    #I'm not sure how to get each \n",
    "    T = launch_and_partials.partials - f_theta\n",
    "\n",
    "    #Includes rearranging the jacobian of profit.\n",
    "\n",
    "    #Return the transitioned values\n",
    "    return  ( A.inverse() ) @  T\n",
    "\n",
    "\n",
    "# This function wraps the single transition and handles updating dates etc.\n",
    "def transition_wrapper(data_in):\n",
    "    \"\"\"\n",
    "    \"\"\"\n",
    "    #unpack states and functions\n",
    "    stocks, debris,profit_fn, laws_motion_fn, neural_network = data_in\n",
    "    \n",
    "    #Calculate new states\n",
    "    launch = neural_network.forward().\n",
    "    new_stocks, new_debris = laws_motion_fn(stocks,debris, launch)\n",
    "\n",
    "    #This gets the transition of the value function derivatives over time.\n",
    "    transitioned = single_transition(\n",
    "                    laws_motion_fn, \n",
    "                    profit_fn, \n",
    "                    stocks, debris, #states\n",
    "                    neural_network #launch function\n",
    "                    )\n",
    "    \n",
    "    #collects the data back together for return, including the updated state variables\n",
    "    data_out = new_stocks, new_debris, profit_fn, laws_motion_fn, transitioned, launch_fn\n",
    "    \n",
    "    return data_out"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "suspected-clerk",
   "metadata": {},
   "source": [
    "## Setup functions related to the problem"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "confused-conclusion",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "    \n",
    "\n",
    "### functions\n",
    "\n",
    "def survival(stock, debris):\n",
    "    \"\"\"\n",
    "    This is a basic, deterministic survival function. \n",
    "    It is based on the CDF of an exponential distribution.\n",
    "    \"\"\"\n",
    "    #SURVIVAL FUNCTION BASED ON AN EXPONENTIAL DISTRIBUTION\n",
    "    return 1 - torch.exp(-SCALING * stock - debris)\n",
    "\n",
    "\n",
    "def laws_of_motion(stock, debris, neural_network):\n",
    "    \"\"\"\n",
    "    This function updates state variables (stock and debris), according \n",
    "    to the laws of motion.\n",
    "    \n",
    "    It returns the state variables as \n",
    "    \"\"\"\n",
    "    launches_and_partials = neural_network.forward(stock,debris)\n",
    "    \n",
    "    s = survival(stock,debris)\n",
    "    #Notes: Survival is a global function.\n",
    "    \n",
    "    new_stock = stock*s + launches_and_partials.launches\n",
    "    \n",
    "\n",
    "    new_debris = (1-DELTA)*debris + LAUNCH_DEBRIS_RATE * launches_and_partials.launches.sum() + COLLISION_DEBRIS_RATE*(1-s) @ 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(stock, debris, launches):\n",
    "    return UTIL_WEIGHTS @ stock - LAUNCH_COST*launches\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "miniature-thread",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Launch functions\n",
    "class ModelMockup():\n",
    "    def __init__(self):\n",
    "        pass\n",
    "        #just a wrapper for NN like results\n",
    "    \n",
    "    def forward(self,stock, debris):\n",
    "        \"\"\"\n",
    "        Temporary launch function \n",
    "        \"\"\"\n",
    "        return LDAPV(torch.ones(len(stocks), requires_grad=True),torch.ones((len(stocks)+len(debris),len(stocks)+len(debris)), requires_grad=True))\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "yellow-frank",
   "metadata": {},
   "source": [
    "# Actual calculations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "enormous-provider",
   "metadata": {},
   "outputs": [],
   "source": [
    "#number of states\n",
    "N = 5\n",
    "\n",
    "#set states\n",
    "stocks = torch.ones(N)\n",
    "#Last one is different if there are more than 1 (just for variety, no theoretical reason)\n",
    "if N > 1:\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",
    "launches = ModelMockup()\n",
    "\n",
    "#Starting point\n",
    "# Stocks, debris, profit fn, laws of motion, item to transition, Launch function\n",
    "base_data = (stocks,debris, profit, laws_of_motion, launches)\n",
    "\n",
    "#Parameters\n",
    "SCALING = torch.ones(5)\n",
    "DELTA = 0.9\n",
    "LAUNCH_DEBRIS_RATE = 0.005\n",
    "LAUNCH_COST = 1.0\n",
    "COLLISION_DEBRIS_RATE = 0.0007\n",
    "UTIL_WEIGHTS = torch.tensor([1,-0.2,0,0,0])\n",
    "BETA = 0.95\n",
    "\n",
    "#Constants determining iterations etc.\n",
    "NUMBER_CONSTELLATIONS = 5\n",
    "NUMBER_DEBRIS_TRACKERS = 1\n",
    "NUMBER_OF_CHOICE_VARIABLES = 1\n",
    "NUMBER_OF_REQUIRED_ITERATED_CONDITIONS = (NUMBER_CONSTELLATIONS+NUMBER_DEBRIS_TRACKERS+(NUMBER_OF_CHOICE_VARIABLES*NUMBER_CONSTELLATIONS))\n",
    "NUMBER_OF_REQUIRED_ITERATIONS = math.ceil(NUMBER_OF_REQUIRED_ITERATED_CONDITIONS/NUMBER_CONSTELLATIONS)\n",
    "NUMBER_OF_REQUIRED_ITERATIONS,NUMBER_OF_REQUIRED_ITERATED_CONDITIONS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "biblical-blake",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Calculate the number of iterations\n",
    "#Get the values from transitions\n",
    "wt1 = compose_recursive_functions(transition_wrapper,NUMBER_OF_REQUIRED_ITERATIONS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "given-clearance",
   "metadata": {},
   "outputs": [],
   "source": [
    "m = ModelMockup()\n",
    "print(m.forward(stocks,debris))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "higher-windsor",
   "metadata": {},
   "source": [
    "# Optimatility conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "breeding-sussex",
   "metadata": {},
   "outputs": [],
   "source": [
    "#Optimality math\n",
    "def optimality(stocks\n",
    "                ,debris\n",
    "                ,profit_fn\n",
    "                ,laws_motion_fn\n",
    "                ,neural_net\n",
    "               ):\n",
    "    \"\"\"\n",
    "    This function takes in the \n",
    "     - stock levels\n",
    "     - debris levels\n",
    "     - profit function\n",
    "     - laws of motion\n",
    "     - results from the neural network\n",
    "    and returns the parts used to make up the optimality conditions\n",
    "    \"\"\"\n",
    "    \n",
    "    #split out the results from the forwarded network\n",
    "    launch_and_partials = neural_net.forward(stock,debris)\n",
    "    \n",
    "    #Derivative of the value function with respect to choice functions\n",
    "    #this returns derivatives with respect to every launch, so I've removed that\n",
    "    fx = jacobian(profit_fn, (stocks,debris,launch_and_partials.launch))[-1][:,0]\n",
    "    \n",
    "    \n",
    "    #The following returns a tuple of tuples of tensors.\n",
    "    #the first tuple contains jacobians related to laws of motion for stocks\n",
    "    #the second tuple contains jacobians related to laws of motion for debris.\n",
    "    #we need the derivatives related to both\n",
    "    b = jacobian(laws_of_motion,(stocks,debris,launch_and_partials.launch))\n",
    "    B = torch.cat((b[0][2],b[1][2].T),axis=1)\n",
    "\n",
    "    #return the optimality values\n",
    "    return fx + BETA *B @ launch_and_partials.partials\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "actual-polish",
   "metadata": {},
   "source": [
    "## Now to set up the recursive set of optimatliy conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "thrown-subject",
   "metadata": {},
   "outputs": [],
   "source": [
    "def recursive_optimality(base_data,transition_wrapper):\n",
    "    #create and return a set of transition wrappers\n",
    "    for f in compose_recursive_functions(transition_wrapper,3):\n",
    "        result = f(base_data)\n",
    "\n",
    "        #unpack results\n",
    "        new_stocks, new_debris, profit_fn, laws_motion_fn, launch_and_partials = result\n",
    "\n",
    "        optimal = optimality(new_stocks\n",
    "                             ,new_debris\n",
    "                             ,profit_fn\n",
    "                             ,laws_motion_fn\n",
    "                             ,launch_and_partials\n",
    "                            )\n",
    "        yield optimal"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "strange-appliance",
   "metadata": {},
   "outputs": [],
   "source": [
    "def optimality_wrapper(stocks,debris, launches):\n",
    "    return optimality(stocks,debris, profit, laws_of_motion, launches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "friendly-acrobat",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimality_wrapper(stocks,debris,launches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "patient-builder",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "markdown",
   "id": "referenced-defense",
   "metadata": {},
   "source": [
    "Notes so far\n",
    " - This takes $\\frac{\\partial W_t}{\\partial{\\theta_t}$ as given. I need to reframe this to clean that out. \n",
    "     - a failed method was just to calculate the partials of W for t = t. That doesn't work.\n",
    "         - Becasue the B value is not invertible, you can't just solve for the partials.\n",
    "         - note the approach below.\n",
    "     - Use part of the launch NN to approximate the gradient conditions.\n",
    "         - Complicates things slightly as it would require adding more outputs to the NN\n",
    "         - Would allow for PDE solution for value function.\n",
    "             - This might be useful for free-entry conditions.\n",
    "             - Not very helpful otherwise.\n",
    "             - I don't think it adds much to the final analysis.\n",
    "     \n",
    "             \n",
    " - There is an issue in how to handle the multiple value functions/loss functions\n",
    "     - Current status\n",
    "         - The launch function (soon to be NN) returns launch choices for each constellation.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "incorrect-carol",
   "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
}