{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "french-experiment",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch.autograd.functional import jacobian\n",
    "import itertools\n",
    "import math\n",
    "import abc\n",
    "\n",
    "class EconomicAgent(metaclass=abc.ABCMeta):\n",
    "    @abc.abstractmethod\n",
    "    def period_benefit(self,state,estimand_interface):\n",
    "        pass\n",
    "    @abc.abstractmethod\n",
    "    def _period_benefit(self):\n",
    "        pass\n",
    "    @abc.abstractmethod\n",
    "    def period_benefit_jacobian_wrt_states(self):\n",
    "        pass\n",
    "    @abc.abstractmethod\n",
    "    def _period_benefit_jacobian_wrt_states(self):\n",
    "        pass\n",
    "    @abc.abstractmethod\n",
    "    def period_benefit_jacobian_wrt_launches(self):\n",
    "        pass\n",
    "    @abc.abstractmethod\n",
    "    def _period_benefit_jacobian_wrt_launches(self):\n",
    "        pass\n",
    "\n",
    "class LinearProfit(EconomicAgent):\n",
    "    \"\"\"\n",
    "    The simplest type of profit function available.\n",
    "    \"\"\"\n",
    "    def __init__(self, constellation_number, discount_factor, benefit_weight, launch_cost, deorbit_cost=0):\n",
    "        #track which constellation this is.\n",
    "        self.constellation_number = constellation_number\n",
    "\n",
    "        #parameters describing the agent's situation\n",
    "        self.discount_factor = discount_factor\n",
    "        self.benefit_weights = benefit_weight\n",
    "        self.launch_cost = launch_cost\n",
    "        self.deorbit_cost = deorbit_cost\n",
    "\n",
    "    def __str__(self):\n",
    "        return \"LinearProfit\\n Benefit weights:\\t{}\\n launch cost:\\t{}\\n Deorbit cost:\\t{}\".format(self.benefit_weights, self.launch_cost, self.deorbit_cost)\n",
    "\n",
    "    def period_benefit(self,state,estimand_interface):\n",
    "        return self._period_benefit(state.stocks, state.debris, estimand_interface.choices)\n",
    "    \n",
    "    def _period_benefit(self,stocks,debris,choice):\n",
    "        profits =  self.benefit_weights @ stocks \\\n",
    "                    - self.launch_cost * choice[self.constellation_number] #\\ \n",
    "                    #- deorbit_cost @ deorbits[self.constellation_number]\n",
    "        return profits\n",
    "\n",
    "    def period_benefit_jacobian_wrt_states(self, states, estimand_interface):\n",
    "        return self._period_benefit_jacobian_wrt_states(states.stocks, states.debris, estimand_interface.choices)\n",
    "\n",
    "    def _period_benefit_jacobian_wrt_states(self, stocks, debris, launches):\n",
    "        jac = jacobian(self._period_benefit, (stocks,debris,launches))\n",
    "        return torch.cat((jac[0], jac[1]))\n",
    "    \n",
    "    def period_benefit_jacobian_wrt_launches(self, states, estimand_interface):\n",
    "        return self._period_benefit_jacobian_wrt_launches(states.stocks, states.debris, estimand_interface.choices)\n",
    "\n",
    "    def _period_benefit_jacobian_wrt_launches(self,stocks,debris,launches):\n",
    "        jac = jacobian(self._period_benefit, (stocks,debris,launches))\n",
    "        return jac[2]\n",
    "\n",
    "class States():\n",
    "    \"\"\"\n",
    "    This is supposed to capture the state variables of the model, to create a common interface \n",
    "    when passing between functions.\n",
    "    \"\"\"\n",
    "    def __init__(self, stocks,debris):\n",
    "        self.stocks = stocks\n",
    "        self.debris = debris\n",
    "        \n",
    "\n",
    "    def __str__(self):\n",
    "        return \"stocks\\t{} \\ndebris\\t {}\".format(self.stocks,self.debris)\n",
    "\n",
    "    @property\n",
    "    def number_constellations(self):\n",
    "        return len(self.stocks)\n",
    "    @property\n",
    "    def number_debris_trackers(self):\n",
    "        return len(self.debris)\n",
    "\n",
    "    \n",
    "class EstimandInterface():\n",
    "    \"\"\"\n",
    "    This defines a clean interface for working with the estimand (i.e. thing we are trying to estimate).\n",
    "    In general, we are trying to estimate the choice variables and the partial derivatives of the value functions.\n",
    "    This \n",
    "\n",
    "    This class wraps output for the neural network (or other estimand), allowing me to \n",
    "        - easily substitute various types of launch functions by having a common interface\n",
    "            - this eases testing\n",
    "        - check dimensionality etc without dealing with randomness\n",
    "            - again, easing testing\n",
    "        - reason more cleanly about the component pieces\n",
    "            - easing programming\n",
    "        - provide a clean interface to find constellation level launch decisions etc.\n",
    "\n",
    "    It takes inputs of two general categories:\n",
    "        - the choice function results\n",
    "        - the partial derivatives of the value function\n",
    "    \"\"\"\n",
    "    def __init__(self, partials, choices, deorbits=None):\n",
    "        self.partials = partials\n",
    "        self.choices = choices\n",
    "        \n",
    "    @property\n",
    "    def number_constellations(self):\n",
    "        pass #fix this\n",
    "        return self.choices.shape[-1]\n",
    "    @property\n",
    "    def number_states(self):\n",
    "        pass #fix this\n",
    "        return self.partials.shape[-1] #This depends on the debris trackers technically.\n",
    "\n",
    "    def choice_single(self, constellation):\n",
    "        #returns the launch decision for the constellation of interest\n",
    "        \n",
    "        filter_tensor = torch.zeros(self.number_constellations)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.choices @ filter_tensor\n",
    "    \n",
    "    def choice_vector(self, constellation):\n",
    "        #returns the launch decision for the constellation of interest as a vector\n",
    "        \n",
    "        filter_tensor = torch.zeros(self.number_constellations)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.choices * filter_tensor\n",
    "    \n",
    "    def partial_vector(self, constellation):\n",
    "        #returns the partials of the value function corresponding to the constellation of interest\n",
    "        \n",
    "        filter_tensor = torch.zeros(self.number_states)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.partials @ filter_tensor\n",
    "    \n",
    "    def partial_matrix(self, constellation):\n",
    "        #returns the partials of the value function corresponding to \n",
    "        #the constellation of interest as a matrix\n",
    "        \n",
    "        filter_tensor = torch.zeros(self.number_states)\n",
    "        filter_tensor[constellation] = 1.0\n",
    "        \n",
    "        return self.partials * filter_tensor\n",
    "    \n",
    "    def __str__(self):\n",
    "        #just a human readable descriptor\n",
    "        return \"Launch Decisions and Partial Derivativs of value function with\\n\\tlaunches\\n\\t\\t {}\\n\\tPartials\\n\\t\\t{}\".format(self.choices,self.partials)\n",
    "\n",
    "\n",
    "class ChoiceFunction(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    This is used to estimate the launch function\n",
    "    \"\"\"\n",
    "    def __init__(self\n",
    "                 ,batch_size\n",
    "                 ,number_states\n",
    "                 ,number_choices\n",
    "                 ,number_constellations\n",
    "                 ,layer_size=12\n",
    "                ):\n",
    "        super().__init__()\n",
    "        \n",
    "        #preprocess\n",
    "        self.preprocess = torch.nn.Linear(in_features=number_states, out_features=layer_size)\n",
    "        \n",
    "        #upsample\n",
    "        self.upsample = lambda x: torch.nn.Upsample(scale_factor=number_constellations)(x).view(batch_size\n",
    "                                                                                            ,number_constellations\n",
    "                                                                                            ,layer_size)\n",
    "        \n",
    "        self.relu = torch.nn.ReLU() #used for coersion to the state space we care about.\n",
    "        \n",
    "       \n",
    "        #sequential steps\n",
    "        self.sequential = torch.nn.Sequential(\n",
    "            torch.nn.Linear(in_features=layer_size, out_features=layer_size)\n",
    "            #who knows if a convolution might help here.\n",
    "            ,torch.nn.Linear(in_features=layer_size, out_features=layer_size)\n",
    "            ,torch.nn.Linear(in_features=layer_size, out_features=layer_size)\n",
    "        )\n",
    "\n",
    "        #reduce the feature axis to match expected results\n",
    "        self.feature_reduction = torch.nn.Linear(in_features=layer_size, out_features=number_choices)\n",
    "\n",
    "        \n",
    "    def forward(self, input_values):\n",
    "        \n",
    "        intermediate_values = self.relu(input_values) #states should be positive anyway.\n",
    "        \n",
    "        intermediate_values = self.preprocess(intermediate_values)\n",
    "        intermediate_values = self.upsample(intermediate_values)\n",
    "        intermediate_values = self.sequential(intermediate_values)\n",
    "        intermediate_values = self.feature_reduction(intermediate_values)\n",
    "        \n",
    "        intermediate_values = self.relu(intermediate_values) #launches are always positive, this may need removed for other types of choices.\n",
    "        \n",
    "        return intermediate_values\n",
    "\n",
    "class PartialDerivativesOfValueEstimand(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    This is used to estimate the partial derivatives of the value functions\n",
    "    \"\"\"\n",
    "    def __init__(self\n",
    "                 ,batch_size\n",
    "                 , number_constellations\n",
    "                 , number_states\n",
    "                 , layer_size=12):\n",
    "        super().__init__()\n",
    "        self.batch_size = batch_size #used for upscaling\n",
    "        self.number_constellations = number_constellations\n",
    "        self.number_states = number_states\n",
    "        self.layer_size = layer_size\n",
    "        \n",
    "        \n",
    "        #preprocess (single linear layer in case there is anything that needs to happen to all states)\n",
    "        self.preprocess = torch.nn.Sequential(\n",
    "            torch.nn.ReLU() #cleanup as states must be positive\n",
    "            ,torch.nn.Linear(in_features = self.number_states, out_features=self.number_states)\n",
    "        )\n",
    "        \n",
    "        #upsample to get the basic dimensionality correct. From (batch,State) to (batch, constellation, state). Includes a reshape\n",
    "        self.upsample = lambda x: torch.nn.Upsample(scale_factor=self.number_constellations)(x).view(self.batch_size\n",
    "                                                                                            ,self.number_constellations\n",
    "                                                                                            ,self.number_states)\n",
    "        \n",
    "        #sequential steps\n",
    "        self.sequential = torch.nn.Sequential(\n",
    "            torch.nn.Linear(in_features=number_states, out_features=layer_size)\n",
    "            #who knows if a convolution or other layer type might help here.\n",
    "            ,torch.nn.Linear(in_features=layer_size, out_features=layer_size)\n",
    "            ,torch.nn.Linear(in_features=layer_size, out_features=layer_size)\n",
    "        )\n",
    "\n",
    "        #reduce the feature axis to match expected results\n",
    "        self.feature_reduction = torch.nn.Linear(in_features=layer_size, out_features=number_states)\n",
    "        \n",
    "    def forward(self, states):\n",
    "        #Note that the input values are just going to be the state variables\n",
    "        #TODO:check that input values match the prepared dimension?\n",
    "        \n",
    "        #preprocess\n",
    "        intermediate = self.preprocess(states)\n",
    "        \n",
    "        #upscale the input values\n",
    "        intermediate = self.upsample(intermediate)\n",
    "        \n",
    "        #intermediate processing\n",
    "        intermediate = self.sequential(intermediate)\n",
    "        \n",
    "        #reduce feature axis to match the expected number of partials\n",
    "        intermediate = self.feature_reduction(intermediate)\n",
    "        \n",
    "        \n",
    "        return intermediate\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "suited-nothing",
   "metadata": {},
   "outputs": [],
   "source": [
    "class EstimandNN(torch.nn.Module):\n",
    "    \"\"\"\n",
    "    This neural network takes the current states as input values and returns both\n",
    "    the partial derivatives of the value function and the launch function.\n",
    "    \"\"\"\n",
    "    def __init__(self\n",
    "                 ,batch_size\n",
    "                 ,number_states\n",
    "                 ,number_choices\n",
    "                 ,number_constellations\n",
    "                 ,layer_size=12\n",
    "                ):\n",
    "        super().__init__()\n",
    "        \n",
    "\n",
    "        self.partials_estimator = PartialDerivativesOfValueEstimand(batch_size, number_constellations, number_states, layer_size)\n",
    "        self.launch_estimator = ChoiceFunction(batch_size, number_states, number_choices, number_constellations, layer_size)\n",
    "        \n",
    "    def forward(self, input_values):\n",
    "        pass\n",
    "        partials = self.partials_estimator(input_values)\n",
    "        launch = self.launch_estimator(input_values)\n",
    "        \n",
    "        return EstimandInterface(partials,launch)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "recognized-story",
   "metadata": {},
   "source": [
    "# Testing\n",
    "\n",
    "Test if states can handle the dimensionality needed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "smart-association",
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size,states,choices = 5,3,1\n",
    "constellations = states -1 #determined by debris tracking\n",
    "max_start_state = 100\n",
    "\n",
    "stocks_and_debris = torch.randint(max_start_state,(batch_size,1,states),dtype=torch.float32)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 84,
   "id": "unsigned-hungary",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "torch.Size([5, 1, 3])"
      ]
     },
     "execution_count": 84,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "stocks_and_debris.size()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "regulated-conversation",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Launch Decisions and Partial Derivativs of value function with\n",
      "\tlaunches\n",
      "\t\t tensor([[[0.0000],\n",
      "         [0.0000]],\n",
      "\n",
      "        [[2.0907],\n",
      "         [0.1053]],\n",
      "\n",
      "        [[2.9730],\n",
      "         [2.2000]],\n",
      "\n",
      "        [[2.3975],\n",
      "         [1.2877]],\n",
      "\n",
      "        [[4.2107],\n",
      "         [2.0752]]], grad_fn=<ReluBackward0>)\n",
      "\tPartials\n",
      "\t\ttensor([[[ 0.1939,  0.3954,  0.0730],\n",
      "         [-0.9428,  0.6145, -0.9247]],\n",
      "\n",
      "        [[ 1.1686,  3.0170,  0.3393],\n",
      "         [-7.1474,  2.3495, -7.0566]],\n",
      "\n",
      "        [[-2.0849,  3.0883, -3.3791],\n",
      "         [-0.6664,  0.0361, -2.2530]],\n",
      "\n",
      "        [[-0.7117,  2.5474, -1.6458],\n",
      "         [-2.1937,  0.6897, -3.0382]],\n",
      "\n",
      "        [[-1.0262,  4.5973, -2.6606],\n",
      "         [-5.4307,  1.4510, -6.6972]]], grad_fn=<AddBackward0>)\n"
     ]
    }
   ],
   "source": [
    "print(a := enn.forward(stocks_and_debris))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "rental-detection",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lossb(a):\n",
    "    #test loss function\n",
    "    return (a**2).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "id": "mechanical-joshua",
   "metadata": {},
   "outputs": [],
   "source": [
    "ch = ChoiceFunction(batch_size\n",
    "                 ,states\n",
    "                 ,choices\n",
    "                 ,constellations\n",
    "                 ,12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "id": "charged-request",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(46.8100, grad_fn=<SumBackward0>)\n",
      "tensor(82442.4219, grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n",
      "tensor(0., grad_fn=<SumBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]],\n",
       "\n",
       "        [[0.],\n",
       "         [0.]]], grad_fn=<ReluBackward0>)"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer = torch.optim.SGD(ch.parameters(),lr=0.01)\n",
    "\n",
    "for i in range(10):\n",
    "    #training loop\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    output = ch.forward(stocks_and_debris)\n",
    "\n",
    "    l = lossb(output)\n",
    "\n",
    "    l.backward()\n",
    "\n",
    "    optimizer.step()\n",
    "\n",
    "    print(l)\n",
    "    \n",
    "\n",
    "ch.forward(stocks_and_debris)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "id": "perceived-permit",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lossc(a):\n",
    "    #test loss function\n",
    "    return (a**2).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 53,
   "id": "atomic-variance",
   "metadata": {},
   "outputs": [],
   "source": [
    "pd = PartialDerivativesOfValueEstimand(\n",
    "                 batch_size\n",
    "                 ,constellations\n",
    "                 ,states\n",
    "                 ,12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "biological-badge",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "tensor(1.9948e-06, grad_fn=<SumBackward0>)\n",
      "tensor(1.7427e-05, grad_fn=<SumBackward0>)\n",
      "tensor(5.7993e-06, grad_fn=<SumBackward0>)\n",
      "tensor(2.9985e-06, grad_fn=<SumBackward0>)\n",
      "tensor(6.5281e-06, grad_fn=<SumBackward0>)\n",
      "tensor(7.8818e-06, grad_fn=<SumBackward0>)\n",
      "tensor(4.4327e-06, grad_fn=<SumBackward0>)\n",
      "tensor(1.1240e-06, grad_fn=<SumBackward0>)\n",
      "tensor(1.2478e-06, grad_fn=<SumBackward0>)\n",
      "tensor(3.5818e-06, grad_fn=<SumBackward0>)\n",
      "tensor(4.3732e-06, grad_fn=<SumBackward0>)\n",
      "tensor(2.7699e-06, grad_fn=<SumBackward0>)\n",
      "tensor(8.9659e-07, grad_fn=<SumBackward0>)\n",
      "tensor(5.7541e-07, grad_fn=<SumBackward0>)\n",
      "tensor(1.5010e-06, grad_fn=<SumBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "tensor([[[ 0.0002, -0.0002, -0.0003],\n",
       "         [ 0.0001, -0.0003, -0.0002]],\n",
       "\n",
       "        [[ 0.0002, -0.0003, -0.0003],\n",
       "         [ 0.0003, -0.0004, -0.0002]],\n",
       "\n",
       "        [[ 0.0002, -0.0003, -0.0003],\n",
       "         [ 0.0002, -0.0003, -0.0003]],\n",
       "\n",
       "        [[ 0.0002, -0.0002, -0.0004],\n",
       "         [ 0.0003, -0.0003, -0.0003]],\n",
       "\n",
       "        [[ 0.0003, -0.0003, -0.0002],\n",
       "         [ 0.0003, -0.0003, -0.0002]]], grad_fn=<AddBackward0>)"
      ]
     },
     "execution_count": 74,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(pd.parameters(),lr=0.0001)\n",
    "\n",
    "for i in range(15):\n",
    "    #training loop\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    output = pd.forward(stocks_and_debris)\n",
    "\n",
    "    l = lossc(output)\n",
    "\n",
    "    l.backward()\n",
    "\n",
    "    optimizer.step()\n",
    "\n",
    "    print(l)\n",
    "    \n",
    "\n",
    "pd.forward(stocks_and_debris)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "compliant-johnson",
   "metadata": {},
   "outputs": [],
   "source": [
    "def lossa(a):\n",
    "    #test loss function\n",
    "    return (a.choices**2).sum() + (a.partials**2).sum()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 81,
   "id": "alive-potato",
   "metadata": {},
   "outputs": [],
   "source": [
    "enn = EstimandNN(batch_size\n",
    "                 ,states\n",
    "                 ,choices\n",
    "                 ,constellations\n",
    "                 ,12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 83,
   "id": "changed-instruction",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 tensor(112.1970, grad_fn=<AddBackward0>)\n",
      "10 tensor(79.8152, grad_fn=<AddBackward0>)\n",
      "20 tensor(55.6422, grad_fn=<AddBackward0>)\n",
      "30 tensor(38.5636, grad_fn=<AddBackward0>)\n",
      "40 tensor(26.9156, grad_fn=<AddBackward0>)\n",
      "50 tensor(18.9986, grad_fn=<AddBackward0>)\n",
      "60 tensor(13.6606, grad_fn=<AddBackward0>)\n",
      "70 tensor(10.1881, grad_fn=<AddBackward0>)\n",
      "80 tensor(8.0395, grad_fn=<AddBackward0>)\n",
      "90 tensor(6.7618, grad_fn=<AddBackward0>)\n",
      "100 tensor(6.0101, grad_fn=<AddBackward0>)\n",
      "110 tensor(5.5517, grad_fn=<AddBackward0>)\n",
      "120 tensor(5.2434, grad_fn=<AddBackward0>)\n",
      "130 tensor(5.0054, grad_fn=<AddBackward0>)\n",
      "140 tensor(4.7988, grad_fn=<AddBackward0>)\n",
      "150 tensor(4.6069, grad_fn=<AddBackward0>)\n",
      "160 tensor(4.4235, grad_fn=<AddBackward0>)\n",
      "170 tensor(4.2468, grad_fn=<AddBackward0>)\n",
      "180 tensor(4.0763, grad_fn=<AddBackward0>)\n",
      "190 tensor(3.9117, grad_fn=<AddBackward0>)\n",
      "200 tensor(3.7532, grad_fn=<AddBackward0>)\n",
      "210 tensor(3.6005, grad_fn=<AddBackward0>)\n",
      "220 tensor(3.4535, grad_fn=<AddBackward0>)\n",
      "230 tensor(3.3121, grad_fn=<AddBackward0>)\n",
      "240 tensor(3.1761, grad_fn=<AddBackward0>)\n",
      "250 tensor(3.0454, grad_fn=<AddBackward0>)\n",
      "260 tensor(2.9198, grad_fn=<AddBackward0>)\n",
      "270 tensor(2.7991, grad_fn=<AddBackward0>)\n",
      "280 tensor(2.6832, grad_fn=<AddBackward0>)\n",
      "290 tensor(2.5720, grad_fn=<AddBackward0>)\n",
      "300 tensor(2.4653, grad_fn=<AddBackward0>)\n",
      "310 tensor(2.3629, grad_fn=<AddBackward0>)\n",
      "320 tensor(2.2646, grad_fn=<AddBackward0>)\n",
      "330 tensor(2.1704, grad_fn=<AddBackward0>)\n",
      "340 tensor(2.0800, grad_fn=<AddBackward0>)\n",
      "350 tensor(1.9933, grad_fn=<AddBackward0>)\n",
      "360 tensor(1.9103, grad_fn=<AddBackward0>)\n",
      "370 tensor(1.8306, grad_fn=<AddBackward0>)\n",
      "380 tensor(1.7543, grad_fn=<AddBackward0>)\n",
      "390 tensor(1.6812, grad_fn=<AddBackward0>)\n",
      "400 tensor(1.6111, grad_fn=<AddBackward0>)\n",
      "410 tensor(1.5440, grad_fn=<AddBackward0>)\n",
      "420 tensor(1.4797, grad_fn=<AddBackward0>)\n",
      "430 tensor(1.4180, grad_fn=<AddBackward0>)\n",
      "440 tensor(1.3590, grad_fn=<AddBackward0>)\n",
      "450 tensor(1.3025, grad_fn=<AddBackward0>)\n",
      "460 tensor(1.2484, grad_fn=<AddBackward0>)\n",
      "470 tensor(1.1965, grad_fn=<AddBackward0>)\n",
      "480 tensor(1.1469, grad_fn=<AddBackward0>)\n",
      "490 tensor(1.0994, grad_fn=<AddBackward0>)\n",
      "500 tensor(1.0540, grad_fn=<AddBackward0>)\n",
      "510 tensor(1.0104, grad_fn=<AddBackward0>)\n",
      "520 tensor(0.9688, grad_fn=<AddBackward0>)\n",
      "530 tensor(0.9290, grad_fn=<AddBackward0>)\n",
      "540 tensor(0.8908, grad_fn=<AddBackward0>)\n",
      "550 tensor(0.8544, grad_fn=<AddBackward0>)\n",
      "560 tensor(0.8195, grad_fn=<AddBackward0>)\n",
      "570 tensor(0.7861, grad_fn=<AddBackward0>)\n",
      "580 tensor(0.7542, grad_fn=<AddBackward0>)\n",
      "590 tensor(0.7237, grad_fn=<AddBackward0>)\n",
      "600 tensor(0.6945, grad_fn=<AddBackward0>)\n",
      "610 tensor(0.6667, grad_fn=<AddBackward0>)\n",
      "620 tensor(0.6400, grad_fn=<AddBackward0>)\n",
      "630 tensor(0.6146, grad_fn=<AddBackward0>)\n",
      "640 tensor(0.5903, grad_fn=<AddBackward0>)\n",
      "650 tensor(0.5671, grad_fn=<AddBackward0>)\n",
      "660 tensor(0.5449, grad_fn=<AddBackward0>)\n",
      "670 tensor(0.5237, grad_fn=<AddBackward0>)\n",
      "680 tensor(0.5035, grad_fn=<AddBackward0>)\n",
      "690 tensor(0.4842, grad_fn=<AddBackward0>)\n",
      "700 tensor(0.4658, grad_fn=<AddBackward0>)\n",
      "710 tensor(0.4482, grad_fn=<AddBackward0>)\n",
      "720 tensor(0.4315, grad_fn=<AddBackward0>)\n",
      "730 tensor(0.4155, grad_fn=<AddBackward0>)\n",
      "740 tensor(0.4002, grad_fn=<AddBackward0>)\n",
      "750 tensor(0.3857, grad_fn=<AddBackward0>)\n",
      "760 tensor(0.3718, grad_fn=<AddBackward0>)\n",
      "770 tensor(0.3586, grad_fn=<AddBackward0>)\n",
      "780 tensor(0.3460, grad_fn=<AddBackward0>)\n",
      "790 tensor(0.3340, grad_fn=<AddBackward0>)\n",
      "800 tensor(0.3226, grad_fn=<AddBackward0>)\n",
      "810 tensor(0.3117, grad_fn=<AddBackward0>)\n",
      "820 tensor(0.3013, grad_fn=<AddBackward0>)\n",
      "830 tensor(0.2914, grad_fn=<AddBackward0>)\n",
      "840 tensor(0.2820, grad_fn=<AddBackward0>)\n",
      "850 tensor(0.2730, grad_fn=<AddBackward0>)\n",
      "860 tensor(0.2645, grad_fn=<AddBackward0>)\n",
      "870 tensor(0.2564, grad_fn=<AddBackward0>)\n",
      "880 tensor(0.2486, grad_fn=<AddBackward0>)\n",
      "890 tensor(0.2413, grad_fn=<AddBackward0>)\n",
      "900 tensor(0.2342, grad_fn=<AddBackward0>)\n",
      "910 tensor(0.2276, grad_fn=<AddBackward0>)\n",
      "920 tensor(0.2212, grad_fn=<AddBackward0>)\n",
      "930 tensor(0.2151, grad_fn=<AddBackward0>)\n",
      "940 tensor(0.2094, grad_fn=<AddBackward0>)\n",
      "950 tensor(0.2039, grad_fn=<AddBackward0>)\n",
      "960 tensor(0.1986, grad_fn=<AddBackward0>)\n",
      "970 tensor(0.1936, grad_fn=<AddBackward0>)\n",
      "980 tensor(0.1889, grad_fn=<AddBackward0>)\n",
      "990 tensor(0.1844, grad_fn=<AddBackward0>)\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "<__main__.EstimandInterface at 0x7f85609fce20>"
      ]
     },
     "execution_count": 83,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "optimizer = torch.optim.Adam(enn.parameters(),lr=0.0001) #note the use of enn in the optimizer\n",
    "\n",
    "for i in range(1000):\n",
    "    #training loop\n",
    "    optimizer.zero_grad()\n",
    "\n",
    "    output = enn.forward(stocks_and_debris)\n",
    "\n",
    "    l = lossa(output)\n",
    "\n",
    "    l.backward()\n",
    "\n",
    "    optimizer.step()\n",
    "\n",
    "    if i%10==0:\n",
    "        print(i, l)\n",
    "    \n",
    "\n",
    "enn.forward(stocks_and_debris)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "proved-amsterdam",
   "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
}